diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index d572107609a3..08341df6c314 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -39,10 +39,18 @@ 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 const& logical_type) { @@ -50,6 +58,12 @@ namespace { 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 get_output_types( std::span 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 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(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>(row_group_indices.size(), std::vector{}); 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>(row_group_indices.size(), std::vector{}); 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>(row_group_indices.size(), std::vector{}); 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>(row_group_indices.size(), std::vector{}); 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; } diff --git a/cpp/tests/io/experimental/hybrid_scan_test.cpp b/cpp/tests/io/experimental/hybrid_scan_test.cpp index f7cd45b926c3..e580994b6d34 100644 --- a/cpp/tests/io/experimental/hybrid_scan_test.cpp +++ b/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{0}; + cudf::test::fixed_width_column_wrapper col0(values, values + num_rows); + cudf::test::fixed_width_column_wrapper 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(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(*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(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(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(row_groups.size())); +} + TEST_F(HybridScanTest, ChunkedReadRowMaskPerPass) { using T = uint32_t;