Hybrid scan reports correct number of input RGs

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

@@ -39,17 +39,31 @@ using parquet::detail::PageNestingDecodeInfo;
 using text::byte_range_info;
 
 namespace {
-// Tests the passed in logical type for a FIXED_LENGTH_BYTE_ARRAY column to see if it should
-// be treated as a string. Currently the only logical type that has special handling is DECIMAL.
-// Other valid types in the future would be UUID (still treated as string) and FLOAT16 (which
-// for now would also be treated as a string).
+
+/**
+ * @brief Tests the logical type for a fixed length byte array column to see if it should be
+ * treated as a string.
+ *
+ * Currently the only logical type that has special handling is DECIMAL. Other valid types in the
+ * future would be UUID (still treated as string) and FLOAT16 (which for now would also be treated
+ * as a string).
+ *
+ * @param logical_type The logical type to test
+ * @return Boolean indicating if the logical type should be treated as a string
+ */
 [[maybe_unused]] inline bool is_treat_fixed_length_as_string(
   cuda::std::optional<LogicalType> const& logical_type)
 {
   if (!logical_type.has_value()) { return true; }
   return logical_type->type != LogicalType::DECIMAL;
 }
 
+/**
+ * @brief Get the output types from the output buffer template
+ *
+ * @param output_buffer_template Output buffer template
+ * @return Output types
+ */
 [[nodiscard]] std::vector<cudf::data_type> get_output_types(
   std::span<inline_column_buffer const> output_buffer_template)
 {
@@ -62,6 +76,22 @@ namespace {
   return output_dtypes;
 }
 
+/**
+ * @brief Count the number of row groups in the input
+ *
+ * @param row_group_indices Row group indices
+ * @return Number of row groups
+ */
+[[nodiscard]] inline size_type count_row_groups(
+  std::span<std::vector<size_type> const> row_group_indices)
+{
+  return std::accumulate(
+    row_group_indices.begin(),
+    row_group_indices.end(),
+    size_type{0},
+    [](auto sum, auto const& rgs) { return sum + static_cast<size_type>(rgs.size()); });
+}
+
 }  // namespace
 
 hybrid_scan_reader_impl::hybrid_scan_reader_impl(
@@ -511,6 +541,8 @@ table_with_metadata hybrid_scan_reader_impl::materialize_filter_columns(
     auto const empty_row_groups =
       std::vector<std::vector<size_type>>(row_group_indices.size(), std::vector<size_type>{});
     prepare_data(read_mode::READ_ALL, empty_row_groups, {}, {});
+    // Set correct number of input row groups to the output metadata
+    _file_itm_data.num_input_row_groups = count_row_groups(row_group_indices);
     return read_chunk_internal(read_mode::READ_ALL, read_columns_mode::FILTER_COLUMNS, row_mask);
   }
 
@@ -547,6 +579,8 @@ table_with_metadata hybrid_scan_reader_impl::materialize_payload_columns(
     auto const empty_row_groups =
       std::vector<std::vector<size_type>>(row_group_indices.size(), std::vector<size_type>{});
     prepare_data(read_mode::READ_ALL, empty_row_groups, {}, {});
+    // Set correct number of input row groups to the output metadata
+    _file_itm_data.num_input_row_groups = count_row_groups(row_group_indices);
     return read_chunk_internal(read_mode::READ_ALL, read_columns_mode::PAYLOAD_COLUMNS, row_mask);
   }
 
@@ -616,6 +650,8 @@ void hybrid_scan_reader_impl::setup_chunking_for_filter_columns(
     auto const empty_row_groups =
       std::vector<std::vector<size_type>>(row_group_indices.size(), std::vector<size_type>{});
     prepare_data(read_mode::CHUNKED_READ, empty_row_groups, {}, {});
+    // Set correct number of input row groups to the output metadata
+    _file_itm_data.num_input_row_groups = count_row_groups(row_group_indices);
     return;
   }
 
@@ -674,6 +710,8 @@ void hybrid_scan_reader_impl::setup_chunking_for_payload_columns(
     auto const empty_row_groups =
       std::vector<std::vector<size_type>>(row_group_indices.size(), std::vector<size_type>{});
     prepare_data(read_mode::CHUNKED_READ, empty_row_groups, {}, {});
+    // Set correct number of input row groups to the output metadata
+    _file_itm_data.num_input_row_groups = count_row_groups(row_group_indices);
     return;
   }
 

cpp/tests/io/experimental/hybrid_scan_test.cpp

@@ -878,6 +878,76 @@ TEST_F(HybridScanTest, StructChildFilterColumn)
     std::invalid_argument);
 }
 
+TEST_F(HybridScanTest, AllRowsPrunedReportsInputRowGroups)
+{
+  using cudf::io::parquet::experimental::use_data_page_mask;
+
+  auto constexpr num_rows = 10;
+
+  // Single row group, single filter column (col0) and single payload column (col1)
+  auto values = cuda::counting_iterator<uint32_t>{0};
+  cudf::test::fixed_width_column_wrapper<uint32_t> col0(values, values + num_rows);
+  cudf::test::fixed_width_column_wrapper<uint32_t> col1(values, values + num_rows);
+  auto table = cudf::table_view{{col0, col1}};
+
+  std::string const filepath = temp_env->get_temp_filepath("AllRowsPruned.parquet");
+  {
+    cudf::io::table_input_metadata input_metadata(table);
+    input_metadata.column_metadata[0].set_name("col0");
+    auto out_opts = cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, table)
+                      .metadata(std::move(input_metadata))
+                      .build();
+    cudf::io::write_parquet(out_opts);
+  }
+
+  auto const stream = cudf::get_default_stream();
+  auto const mr     = cudf::get_current_device_resource_ref();
+
+  auto scalar    = cudf::numeric_scalar<uint32_t>(0, true, stream);
+  auto literal   = cudf::ast::literal(scalar);
+  auto col_ref_0 = cudf::ast::column_name_reference("col0");
+  auto filter_expression =
+    cudf::ast::operation(cudf::ast::ast_operator::GREATER_EQUAL, col_ref_0, literal);
+
+  auto options = cudf::io::parquet_reader_options::builder().filter(filter_expression).build();
+
+  auto datasource          = cudf::io::datasource::create(filepath);
+  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 row_groups = reader->all_row_groups(options);
+  auto false_iter       = cuda::make_constant_iterator(false);
+  auto row_mask =
+    cudf::test::fixed_width_column_wrapper<bool>(false_iter, false_iter + num_rows).release();
+  auto row_mask_view = row_mask->mutable_view();
+
+  auto const filter_byte_ranges = reader->filter_column_chunks_byte_ranges(row_groups, options);
+  auto [filter_buffers, filter_data, filter_tasks] =
+    cudf::io::parquet::fetch_byte_ranges_to_device_async(
+      *datasource, filter_byte_ranges, stream, mr);
+  filter_tasks.get();
+
+  auto const filter_result = reader->materialize_filter_columns(
+    row_groups, filter_data, row_mask_view, use_data_page_mask::YES, options, stream, mr);
+  EXPECT_EQ(filter_result.tbl->num_rows(), 0);
+  EXPECT_EQ(filter_result.metadata.num_input_row_groups,
+            static_cast<cudf::size_type>(row_groups.size()));
+
+  // Payload columns: same expectation
+  auto const payload_byte_ranges = reader->payload_column_chunks_byte_ranges(row_groups, options);
+  auto [payload_buffers, payload_data, payload_tasks] =
+    cudf::io::parquet::fetch_byte_ranges_to_device_async(
+      *datasource, payload_byte_ranges, stream, mr);
+  payload_tasks.get();
+
+  auto const payload_result = reader->materialize_payload_columns(
+    row_groups, payload_data, row_mask_view, use_data_page_mask::YES, options, stream, mr);
+  EXPECT_EQ(payload_result.tbl->num_rows(), 0);
+  EXPECT_EQ(payload_result.metadata.num_input_row_groups,
+            static_cast<cudf::size_type>(row_groups.size()));
+}
+
 TEST_F(HybridScanTest, ChunkedReadRowMaskPerPass)
 {
   using T = uint32_t;