diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index dd44c53e5494..aab25c0e648b 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -793,38 +793,28 @@ hybrid_scan_reader_impl::construct_row_group_passes( CUDF_EXPECTS( pass_read_limit > 0, "Pass read limit must be greater than 0", std::invalid_argument); - auto row_groups_info = std::vector{}; - row_groups_info.reserve(total_row_groups); - size_t start_row = 0; + auto row_group_ids = std::vector>{}; + auto row_group_sizes = std::vector{}; + row_group_ids.reserve(total_row_groups); + row_group_sizes.reserve(total_row_groups); + std::for_each(cuda::counting_iterator(0), cuda::counting_iterator(row_group_indices.size()), [&](auto const source_index) { - auto const& src_row_groups = row_group_indices[source_index]; - std::transform( - src_row_groups.begin(), - src_row_groups.end(), - std::back_inserter(row_groups_info), - [&](auto const rg_index) { - auto const& row_group = - _extended_metadata->get_row_group(rg_index, source_index); - auto const [compressed_size, total_size, num_rows, max_leaf_values] = - _extended_metadata->get_row_group_properties(row_group); - auto rg_info = row_group_info{.index = rg_index, - .start_row = start_row, - .unadjusted_num_rows = num_rows, - .source_index = source_index, - .compressed_size = compressed_size, - .max_leaf_values = max_leaf_values}; - start_row += num_rows; - return rg_info; - }); + for (auto const rg_index : row_group_indices[source_index]) { + row_group_ids.emplace_back(rg_index, source_index); + // TODO(mh): Compute the row group size information over the selected columns + // instead + row_group_sizes.push_back(_extended_metadata->get_row_group_size_info( + rg_index, source_index, std::nullopt)); + } }); auto const comp_read_limit = static_cast( pass_read_limit * cudf::io::parquet::detail::input_limit_compression_reserve); auto const pass_data = - cudf::io::parquet::detail::compute_row_group_passes(row_groups_info, comp_read_limit, 0); + cudf::io::parquet::detail::compute_row_group_passes(row_group_sizes, comp_read_limit, 0); // Convert offset-based pass boundaries back to vectors of row group indices auto const& offsets = pass_data.pass_row_group_offsets; @@ -832,7 +822,7 @@ hybrid_scan_reader_impl::construct_row_group_passes( passes.reserve(offsets.size() - 1); auto row_group_source_map = std::vector{}; auto const has_multiple_sources = row_group_indices.size() > 1; - if (has_multiple_sources) { row_group_source_map.reserve(row_groups_info.size()); } + if (has_multiple_sources) { row_group_source_map.reserve(row_group_ids.size()); } std::transform(offsets.begin(), offsets.end() - 1, offsets.begin() + 1, @@ -840,12 +830,12 @@ hybrid_scan_reader_impl::construct_row_group_passes( [&](auto const start, auto const end) { auto pass = std::vector{}; pass.reserve(end - start); - std::for_each(row_groups_info.begin() + start, - row_groups_info.begin() + end, - [&](auto const& rg_info) { - pass.emplace_back(rg_info.index); + std::for_each(row_group_ids.begin() + start, + row_group_ids.begin() + end, + [&](auto const& row_group_id) { + pass.emplace_back(row_group_id.first); if (has_multiple_sources) { - row_group_source_map.emplace_back(rg_info.source_index); + row_group_source_map.emplace_back(row_group_id.second); } }); return pass; diff --git a/cpp/src/io/parquet/reader_impl_chunking.cu b/cpp/src/io/parquet/reader_impl_chunking.cu index c2a53826bb05..6f1c9c05e3cf 100644 --- a/cpp/src/io/parquet/reader_impl_chunking.cu +++ b/cpp/src/io/parquet/reader_impl_chunking.cu @@ -543,8 +543,19 @@ void reader_impl::compute_input_passes(read_mode mode) ? static_cast(_input_pass_read_limit * input_limit_compression_reserve) : std::numeric_limits::max(); + // Compute size information for each row group by the columns we are actually going to read. + auto row_group_sizes = std::vector{}; + row_group_sizes.reserve(row_groups_info.size()); + std::transform(row_groups_info.cbegin(), + row_groups_info.cend(), + std::back_inserter(row_group_sizes), + [&](auto const& row_group) { + return _metadata->get_row_group_size_info( + row_group.index, row_group.source_index, _input_columns); + }); + auto pass_data = - compute_row_group_passes(row_groups_info, comp_read_limit, _file_itm_data.global_skip_rows); + compute_row_group_passes(row_group_sizes, comp_read_limit, _file_itm_data.global_skip_rows); _file_itm_data.input_pass_row_group_offsets = std::move(pass_data.pass_row_group_offsets); _file_itm_data.input_pass_start_row_count = std::move(pass_data.pass_start_row_counts); diff --git a/cpp/src/io/parquet/reader_impl_chunking_utils.cu b/cpp/src/io/parquet/reader_impl_chunking_utils.cu index ebca8918171c..6017d7bfa426 100644 --- a/cpp/src/io/parquet/reader_impl_chunking_utils.cu +++ b/cpp/src/io/parquet/reader_impl_chunking_utils.cu @@ -953,7 +953,7 @@ rmm::device_uvector compute_level_decode_sizes(device_span row_groups_info, +row_group_pass_data compute_row_group_passes(std::span row_group_sizes, std::size_t comp_read_limit, int64_t skip_rows) { @@ -970,14 +970,18 @@ row_group_pass_data compute_row_group_passes(cudf::host_span 1) - ? (rgi.start_row + rgi.unadjusted_num_rows - skip_rows) - : rgi.unadjusted_num_rows; + auto row_group_rows = rgi.unadjusted_num_rows; + if (row_group_sizes.size() > 1) { + CUDF_EXPECTS(std::cmp_greater_equal(rgi.unadjusted_num_rows, skip_rows), + "Row groups must contribute non-negative effective rows", + std::invalid_argument); + row_group_rows -= skip_rows; + } auto const compressed_rg_size = rgi.compressed_size; auto const row_group_leaf_values = rgi.max_leaf_values; @@ -1026,8 +1030,8 @@ row_group_pass_data compute_row_group_passes(cudf::host_span row_groups_info, +row_group_pass_data compute_row_group_passes(std::span row_group_sizes, std::size_t comp_read_limit, int64_t skip_rows); diff --git a/cpp/src/io/parquet/reader_impl_helpers.cpp b/cpp/src/io/parquet/reader_impl_helpers.cpp index 6639fd03e47c..4d1ecc404376 100644 --- a/cpp/src/io/parquet/reader_impl_helpers.cpp +++ b/cpp/src/io/parquet/reader_impl_helpers.cpp @@ -1236,6 +1236,49 @@ RowGroup const& aggregate_reader_metadata::get_row_group(size_type row_group_ind return per_file_metadata[src_idx].row_groups[row_group_index]; } +row_group_size_info aggregate_reader_metadata::get_row_group_size_info( + size_type row_group_index, + size_type src_idx, + std::optional> input_columns) const +{ + auto const& row_group = get_row_group(row_group_index, src_idx); + + CUDF_EXPECTS(row_group.num_rows >= 0 && std::in_range(row_group.num_rows), + "Row group has an invalid number of rows", + std::invalid_argument); + auto size_info = + row_group_size_info{.unadjusted_num_rows = static_cast(row_group.num_rows)}; + + // Helper function to overflow-safeadd compressed sizes + auto const add_compressed_size = [&](auto&& current_size, auto&& colchunk_compressed_size) { + auto const sum = cuda::add_overflow(current_size, colchunk_compressed_size); + CUDF_EXPECTS(not sum.overflow, + "Row group compressed size exceeds the supported range", + std::overflow_error); + return sum.value; + }; + + if (input_columns.has_value()) { + for (auto const& column : *input_columns) { + auto const& column_metadata = + get_column_metadata(row_group_index, src_idx, column.schema_idx); + size_info.compressed_size = + add_compressed_size(size_info.compressed_size, column_metadata.total_compressed_size); + size_info.max_leaf_values = + std::max(size_info.max_leaf_values, column_metadata.num_values); + } + } else { + for (auto const& column_chunk : row_group.columns) { + size_info.compressed_size = add_compressed_size(size_info.compressed_size, + column_chunk.meta_data.total_compressed_size); + size_info.max_leaf_values = + std::max(size_info.max_leaf_values, column_chunk.meta_data.num_values); + } + } + + return size_info; +} + ColumnChunkMetaData const& aggregate_reader_metadata::get_column_metadata(size_type row_group_index, size_type src_idx, int schema_idx) const @@ -1412,31 +1455,6 @@ std::vector aggregate_reader_metadata::get_pandas_index_names() con return names; } -std::tuple aggregate_reader_metadata::get_row_group_properties( - RowGroup const& row_group) const -{ - auto const compressed_size = std::transform_reduce( - row_group.columns.cbegin(), - row_group.columns.cend(), - size_t{0}, - std::plus<>(), - [](auto const& colchunk) { return colchunk.meta_data.total_compressed_size; }); - - auto const total_size = compressed_size + row_group.total_byte_size; - - size_t const max_leaf_values = - row_group.columns.empty() - ? 0 - : std::max_element(row_group.columns.cbegin(), - row_group.columns.cend(), - [](auto const& a, auto const& b) { - return a.meta_data.num_values < b.meta_data.num_values; - }) - ->meta_data.num_values; - - return {compressed_size, total_size, static_cast(row_group.num_rows), max_leaf_values}; -} - std::tuple>, std::vector>, @@ -1790,10 +1808,6 @@ aggregate_reader_metadata::select_row_groups( // Update the number of rows read from this data source num_rows_per_source[src_idx] += num_rows_this_row_group; - // Get row group properties - auto const [compressed_size, total_size, num_rows, max_leaf_values] = - get_row_group_properties(rg); - // We need the unadjusted start index of this row group to correctly // initialize ColumnChunkDesc for this row group in // create_global_chunk_info() and calculate the row offset for the first @@ -1802,10 +1816,8 @@ aggregate_reader_metadata::select_row_groups( 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(src_idx), - .compressed_size = compressed_size, - .max_leaf_values = max_leaf_values}); + .unadjusted_num_rows = static_cast(rg.num_rows), + .source_index = static_cast(src_idx)}); // If page-level indexes are present, then collect extra chunk and page // info. The page indexes rely on absolute row numbers - not adjusted for diff --git a/cpp/src/io/parquet/reader_impl_helpers.hpp b/cpp/src/io/parquet/reader_impl_helpers.hpp index b275d8ff678b..ed67ca05cfc3 100644 --- a/cpp/src/io/parquet/reader_impl_helpers.hpp +++ b/cpp/src/io/parquet/reader_impl_helpers.hpp @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -66,19 +67,26 @@ struct column_chunk_info { * @brief The row_group_info class */ 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_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; // global start row of this row group 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 - size_t max_leaf_values; // maximum number of leaf values in the row group // Optional metadata pulled from the column and offset indexes, if present. std::optional> column_chunks; }; +/** + * @brief Row group size information for pass partitioning. + */ +struct row_group_size_info { + size_t unadjusted_num_rows; // number of unadjusted rows in this row group + size_t compressed_size; // compressed size of the selected columns in this row group + size_t max_leaf_values; // maximum number of leaf values over the selected columns +}; + /** * @brief Translates Parquet datatype to cuDF type enum */ @@ -434,6 +442,22 @@ class aggregate_reader_metadata { */ [[nodiscard]] RowGroup const& get_row_group(size_type row_group_index, size_type src_idx) const; + /** + * @brief Computes row group size information over selected columns + * + * When `input_columns` is specified, computes the compressed size and maximum leaf value count + * over only those columns. Otherwise, over all columns in the row group. + * + * @param row_group_index Index of the row group within its source + * @param src_idx Index of the input source + * @param input_columns Optional selected leaf columns + * @return Row group size information + */ + [[nodiscard]] row_group_size_info get_row_group_size_info( + size_type row_group_index, + size_type src_idx, + std::optional> input_columns) const; + /** * @brief Check if all row groups have an offset index * @@ -622,18 +646,6 @@ class aggregate_reader_metadata { */ [[nodiscard]] std::vector get_pandas_index_names() const; - /** - * @brief Computes the compressed and total size, the number of rows, and the maximum number of - * leaf values in the specified row group - * - * @param row_group The row group - * - * @return A tuple of row group compressed size, total size, number of rows, and maximum leaf - * values - */ - [[nodiscard]] std::tuple get_row_group_properties( - RowGroup const& rg) const; - /** * @brief Filters the row groups using stats and bloom filters based on predicate filter * diff --git a/cpp/tests/io/parquet_chunked_reader_test.cu b/cpp/tests/io/parquet_chunked_reader_test.cu index c131a9983e01..4e259ab719d1 100644 --- a/cpp/tests/io/parquet_chunked_reader_test.cu +++ b/cpp/tests/io/parquet_chunked_reader_test.cu @@ -103,14 +103,8 @@ auto write_file(std::vector>& input_columns, return std::pair{std::move(input_table), std::move(filepath)}; } -auto chunked_read(std::vector const& filepaths, - std::size_t output_limit, - std::size_t input_limit = 0) +auto chunked_read(cudf::io::chunked_parquet_reader const& reader) { - auto const read_opts = - cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepaths}).build(); - auto reader = cudf::io::chunked_parquet_reader(output_limit, input_limit, read_opts); - auto num_chunks = 0; auto out_tables = std::vector>{}; @@ -133,6 +127,23 @@ auto chunked_read(std::vector const& filepaths, return std::pair(cudf::concatenate(out_tviews), num_chunks); } +auto chunked_read(cudf::io::parquet_reader_options const& read_opts, + std::size_t output_limit, + std::size_t input_limit = 0) +{ + auto reader = cudf::io::chunked_parquet_reader(output_limit, input_limit, read_opts); + return chunked_read(reader); +} + +auto chunked_read(std::vector const& filepaths, + std::size_t output_limit, + std::size_t input_limit = 0) +{ + auto const read_opts = + cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepaths}).build(); + return chunked_read(read_opts, output_limit, input_limit); +} + auto chunked_read(std::string const& filepath, std::size_t output_limit, std::size_t input_limit = 0) @@ -149,27 +160,7 @@ auto chunked_read(std::vector>&& sources, auto const read_opts = cudf::io::parquet_reader_options::builder().build(); auto reader = cudf::io::chunked_parquet_reader( output_limit, input_limit, std::move(sources), std::move(metadatas), read_opts); - - auto num_chunks = 0; - auto out_tables = std::vector>{}; - - do { - auto chunk = reader.read_chunk(); - // If the input file is empty, the first call to `read_chunk` will return an empty table. - // Thus, we only check for non-empty output table from the second call. - if (num_chunks > 0) { - CUDF_EXPECTS(chunk.tbl->num_rows() != 0, "Number of rows in the new chunk is zero."); - } - ++num_chunks; - out_tables.emplace_back(std::move(chunk.tbl)); - } while (reader.has_next()); - - auto out_tviews = std::vector{}; - for (auto const& tbl : out_tables) { - out_tviews.emplace_back(tbl->view()); - } - - return std::pair(cudf::concatenate(out_tviews), num_chunks); + return chunked_read(reader); } auto const read_table_and_nrows_per_source(cudf::io::chunked_parquet_reader const& reader) @@ -1338,6 +1329,55 @@ TEST_F(ParquetChunkedReaderInputLimitConstrainedTest, MixedColumns) struct ParquetChunkedReaderInputLimitTest : public cudf::test::BaseFixture {}; +TEST_F(ParquetChunkedReaderInputLimitTest, ProjectedColumnsReducePasses) +{ + constexpr int num_columns = 16; + constexpr int num_rows = 2'000; + constexpr int rows_per_row_group = num_rows / 2; + auto const filepath = temp_env->get_temp_filepath("ProjectedColumnPasses.parquet"); + + // Input table + auto columns = std::vector>{}; + std::transform(cuda::counting_iterator{0}, + cuda::counting_iterator{num_columns}, + std::back_inserter(columns), + [](int c) { + return cudf::test::fixed_width_column_wrapper( + cuda::counting_iterator{0}, cuda::counting_iterator{num_rows}) + .release(); + }); + auto const input_table = cudf::table(std::move(columns)); + + // Write table to parquet + cudf::io::write_parquet( + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, input_table.view()) + .metadata(cudf::io::table_input_metadata{input_table.view()}) + .compression(cudf::io::compression_type::NONE) + .dictionary_policy(cudf::io::dictionary_policy::NEVER) + .max_page_fragment_size(rows_per_row_group) + .row_group_size_rows(rows_per_row_group) + .build()); + + // Each row group is roughly 64 KB = sizeof(int32) * 1000 * 16. Use 100KB as limit so we pick one + // row group per pass when all columns are read but more than one for single column read. + constexpr std::size_t pass_read_limit = 100'000; + + auto const all_columns = + cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}).build(); + auto const [all_cols, all_cols_chunks] = chunked_read(all_columns, 0, pass_read_limit); + + auto const one_column = cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}) + .column_indices({0}) + .build(); + auto const [one_col, one_col_chunks] = chunked_read(one_column, 0, pass_read_limit); + + CUDF_TEST_EXPECT_TABLES_EQUAL(all_cols->view(), input_table.view()); + CUDF_TEST_EXPECT_TABLES_EQUAL(one_col->view(), cudf::table_view{{input_table.view().column(0)}}); + + // Projected read must yield fewer chunks than full read + EXPECT_LT(one_col_chunks, all_cols_chunks); +} + namespace { struct offset_gen { int const group_size;