From fc55575b37e74006039172db200bc84f7610a49b Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Thu, 2 Jul 2026 01:48:00 +0000 Subject: [PATCH 1/7] Prepend file-local row index column in Parquet reader --- cpp/include/cudf/io/parquet.hpp | 35 ++++ cpp/src/io/parquet/reader_impl.cpp | 30 +++- cpp/src/io/parquet/reader_impl.hpp | 15 ++ cpp/src/io/parquet/reader_impl_helpers.cpp | 21 +++ cpp/src/io/parquet/reader_impl_helpers.hpp | 9 + cpp/src/io/parquet/reader_impl_preprocess.cu | 81 +++++++++ cpp/tests/io/parquet_chunked_reader_test.cu | 143 ++++++++++++++++ cpp/tests/io/parquet_reader_test.cpp | 167 +++++++++++++++++++ 8 files changed, 494 insertions(+), 7 deletions(-) diff --git a/cpp/include/cudf/io/parquet.hpp b/cpp/include/cudf/io/parquet.hpp index 0c895abfa0a1..a25eca1ddbfd 100644 --- a/cpp/include/cudf/io/parquet.hpp +++ b/cpp/include/cudf/io/parquet.hpp @@ -108,6 +108,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> _reader_column_schema; @@ -310,6 +312,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 * @@ -561,6 +577,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; } }; /** @@ -834,6 +857,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. */ diff --git a/cpp/src/io/parquet/reader_impl.cpp b/cpp/src/io/parquet/reader_impl.cpp index aef720839fa2..ee123667f56c 100644 --- a/cpp/src/io/parquet/reader_impl.cpp +++ b/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(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,23 @@ 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) { + prepend_row_index_column(read_info, out_columns); + out_metadata.schema_info.emplace(out_metadata.schema_info.begin(), + column_name_info{.name = "row_idx", .is_nullable = false}); + } + 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}); + } } // Offset column references in `_expr_conv` by the number of prepended columns - auto const num_prepended_cols = static_cast(_options.prepend_source_index_column); + auto const num_prepended_cols = static_cast(_options.prepend_source_index_column) + + static_cast(_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(); diff --git a/cpp/src/io/parquet/reader_impl.hpp b/cpp/src/io/parquet/reader_impl.hpp index 060c22d82602..ef16c3487eef 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -429,6 +429,19 @@ class reader_impl { void prepend_source_index_column(std::span num_rows_per_source, std::vector>& out_columns); + /** + * @brief Construct and prepend the file-local row index column to the output columns + * + * 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 + * @param out_columns Current output columns + */ + void prepend_row_index_column(row_range const& read_info, + std::vector>& out_columns); + /** * @brief Computes the names of columns to be read from the file, if specified. * @@ -469,6 +482,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 diff --git a/cpp/src/io/parquet/reader_impl_helpers.cpp b/cpp/src/io/parquet/reader_impl_helpers.cpp index 18091cf35822..8925e38a0ac5 100644 --- a/cpp/src/io/parquet/reader_impl_helpers.cpp +++ b/cpp/src/io/parquet/reader_impl_helpers.cpp @@ -558,6 +558,22 @@ std::vector aggregate_reader_metadata::get_num_row_groups_per_file() return per_file_num_row_groups; } +std::vector 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 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{}, + [](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(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(src_idx), .compressed_size = compressed_size, diff --git a/cpp/src/io/parquet/reader_impl_helpers.hpp b/cpp/src/io/parquet/reader_impl_helpers.hpp index 9417416c1e57..631293b3789e 100644 --- a/cpp/src/io/parquet/reader_impl_helpers.hpp +++ b/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 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 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 * diff --git a/cpp/src/io/parquet/reader_impl_preprocess.cu b/cpp/src/io/parquet/reader_impl_preprocess.cu index 2b9723c03b63..721f1c30f9fe 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess.cu @@ -22,7 +22,10 @@ #include +#include #include +#include +#include #include #include #include @@ -1125,6 +1128,84 @@ cudf::detail::host_vector reader_impl::calculate_page_string_offsets() return cudf::detail::make_pinned_vector(d_col_sizes, _stream); } +namespace { + +/** + * @brief Maps each global row index in the current row range to the source-local row index + */ +struct map_global_to_local_row_index { + std::size_t const* rg_global_row_offsets; ///< Global row offsets for each selected row group + std::size_t const* rg_local_row_offsets; ///< File-local row offsets for each selected row group + std::size_t num_row_groups; + + __device__ std::size_t operator()(std::size_t global_row_index) const + { + auto const rg_idx = thrust::upper_bound(thrust::seq, + rg_global_row_offsets, + rg_global_row_offsets + num_row_groups, + global_row_index) - + rg_global_row_offsets - 1; + return global_row_index + rg_local_row_offsets[rg_idx]; + } +}; + +} // namespace + +void reader_impl::prepend_row_index_column(row_range const& read_info, + std::vector>& out_columns) +{ + if (not _options.prepend_row_index_column) { return; } + + using column_type = size_t; + auto constexpr dtype = cudf::data_type{cudf::type_to_id()}; + + auto const prepend_column = [&](auto&& column) { + out_columns.emplace(out_columns.begin(), std::move(column)); + }; + + // Empty column + if (read_info.num_rows == 0) { + prepend_column(cudf::make_empty_column(dtype)); + return; + } + + // Collect the global and file-local row offsets for each currently selected row group (in + // selection order). + auto const& row_groups = _file_itm_data.row_groups; + auto host_rg_global_offsets = + cudf::detail::make_empty_pinned_vector(row_groups.size(), _stream); + auto host_rg_local_offsets = + cudf::detail::make_empty_pinned_vector(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 - rg.start_row); + } + + // Map global row indices in the current row-range to corresponding source-local row indices + rmm::device_uvector col_data(read_info.num_rows, _stream, _mr); + { + auto temp_mr = cudf::get_current_device_resource_ref(); + + // Copy row offsets to device + auto const rg_global_offsets = + cudf::detail::make_device_uvector_async(host_rg_global_offsets, _stream, temp_mr); + auto const rg_local_offsets = + cudf::detail::make_device_uvector_async(host_rg_local_offsets, _stream, temp_mr); + + // 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(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(); + } + + prepend_column(std::make_unique(std::move(col_data), rmm::device_buffer{}, 0)); +} + void reader_impl::prepend_source_index_column(std::span num_rows_per_source, std::vector>& out_columns) { diff --git a/cpp/tests/io/parquet_chunked_reader_test.cu b/cpp/tests/io/parquet_chunked_reader_test.cu index d9eb03a867d6..9edd60d471a3 100644 --- a/cpp/tests/io/parquet_chunked_reader_test.cu +++ b/cpp/tests/io/parquet_chunked_reader_test.cu @@ -1983,6 +1983,149 @@ TEST_F(ParquetChunkedReaderTest, TestNumRowsPerSourceMultipleSources) } } +TEST_F(ParquetChunkedReaderTest, TestRowIndexColumnMultipleSources) +{ + constexpr int num_rows = 10'723; // A prime number + constexpr int rows_in_row_group = 500; + + // Table with single col of random int64 values + auto const int64_data = random_values(num_rows); + auto int64_col = int64s_col(int64_data.begin(), int64_data.end()).release(); + + std::vector> input_columns; + input_columns.emplace_back(std::move(int64_col)); + + // Write to Parquet + auto const [expected, filepath] = write_file(input_columns, + "row_index_column", + false, + false, + cudf::io::default_max_page_size_bytes, + rows_in_row_group); + + // Chunked-read four data sources entirely with prepended source and row index columns + { + auto const nsources = 4; + auto constexpr output_read_limit = 15'000; + auto constexpr pass_read_limit = 35'000; + std::vector const datasources(nsources, filepath); + + auto const options = + cudf::io::parquet_reader_options_builder(cudf::io::source_info{datasources}) + .prepend_source_index_column(true) + .prepend_row_index_column(true) + .build(); + auto const reader = cudf::io::chunked_parquet_reader( + output_read_limit, pass_read_limit, options, cudf::get_default_stream()); + + auto const [result, num_chunks, num_rows_per_source] = read_table_and_nrows_per_source(reader); + + EXPECT_EQ(result->num_columns(), 3); + + // Expected source and (file-local) row index columns + auto const src_index = cudf::detail::make_counting_transform_iterator( + 0, [](cudf::size_type i) { return i / num_rows; }); + auto const expected_src_index = cudf::test::fixed_width_column_wrapper( + src_index, src_index + nsources * num_rows); + auto const row_index = cudf::detail::make_counting_transform_iterator( + 0, [](cudf::size_type i) -> size_t { return i % num_rows; }); + auto const expected_row_index = + cudf::test::fixed_width_column_wrapper(row_index, row_index + nsources * num_rows); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view().column(0), expected_src_index); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view().column(1), expected_row_index); + } + + // Chunked-read rows_to_read rows skipping rows_to_skip from four data sources + { + auto const rows_to_skip = 15'201; + auto const rows_to_read = 20'232; + auto constexpr output_read_limit = 15'000; + auto constexpr pass_read_limit = 35'000; + auto const nsources = 4; + std::vector const datasources(nsources, filepath); + + auto const options = + cudf::io::parquet_reader_options_builder(cudf::io::source_info{datasources}) + .skip_rows(rows_to_skip) + .num_rows(rows_to_read) + .prepend_source_index_column(true) + .prepend_row_index_column(true) + .build(); + auto const reader = cudf::io::chunked_parquet_reader( + output_read_limit, pass_read_limit, options, cudf::get_default_stream()); + + auto const [result, num_chunks, num_rows_per_source] = read_table_and_nrows_per_source(reader); + + EXPECT_EQ(result->num_columns(), 3); + EXPECT_EQ(result->num_rows(), rows_to_read); + + // Expected source and (file-local) row index columns + auto const src_index = cudf::detail::make_counting_transform_iterator( + rows_to_skip, [](cudf::size_type i) { return i / num_rows; }); + auto const expected_src_index = + cudf::test::fixed_width_column_wrapper(src_index, src_index + rows_to_read); + auto const row_index = cudf::detail::make_counting_transform_iterator( + rows_to_skip, [](cudf::size_type i) -> size_t { return i % num_rows; }); + auto const expected_row_index = + cudf::test::fixed_width_column_wrapper(row_index, row_index + rows_to_read); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view().column(0), expected_src_index); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view().column(1), expected_row_index); + } + + // Filtered chunked-read from four data sources with prepended source and row index columns. The + // row index must reflect each surviving row's original file-local position, not its output row. + { + auto const nsources = 4; + auto constexpr output_read_limit = 15'000; + auto constexpr pass_read_limit = 35'000; + auto const max_value = int64_data[int64_data.size() / 2]; + std::vector const datasources(nsources, filepath); + + auto literal_value = cudf::numeric_scalar{max_value}; + auto literal = cudf::ast::literal{literal_value}; + auto col_ref = cudf::ast::column_reference(0); + auto filter = cudf::ast::operation(cudf::ast::ast_operator::LESS_EQUAL, col_ref, literal); + + auto const options = + cudf::io::parquet_reader_options_builder(cudf::io::source_info{datasources}) + .filter(filter) + .prepend_source_index_column(true) + .prepend_row_index_column(true) + .build(); + auto const reader = cudf::io::chunked_parquet_reader( + output_read_limit, pass_read_limit, options, cudf::get_default_stream()); + + auto const [result, num_chunks, num_rows_per_source] = read_table_and_nrows_per_source(reader); + + EXPECT_EQ(result->num_columns(), 3); + + // Per source: the file-local indices (and values) of the rows that survive the filter + std::vector src_index_data; + std::vector row_index_data; + std::vector values_data; + for (cudf::size_type s = 0; s < nsources; ++s) { + for (cudf::size_type i = 0; i < num_rows; ++i) { + if (int64_data[i] <= max_value) { + src_index_data.push_back(s); + row_index_data.push_back(static_cast(i)); + values_data.push_back(int64_data[i]); + } + } + } + + auto const expected_src_index = + cudf::test::fixed_width_column_wrapper(src_index_data.begin(), + src_index_data.end()); + auto const expected_row_index = + cudf::test::fixed_width_column_wrapper(row_index_data.begin(), row_index_data.end()); + auto const expected_values = + int64s_col(values_data.begin(), values_data.end()); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view().column(0), expected_src_index); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view().column(1), expected_row_index); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view().column(2), expected_values); + } +} + TEST_F(ParquetChunkedReaderTest, TestNumRowsPerSourceEmptyTable) { auto constexpr output_read_limit = 4'500; diff --git a/cpp/tests/io/parquet_reader_test.cpp b/cpp/tests/io/parquet_reader_test.cpp index b28ec897c956..d08ec62bac37 100644 --- a/cpp/tests/io/parquet_reader_test.cpp +++ b/cpp/tests/io/parquet_reader_test.cpp @@ -4977,3 +4977,170 @@ TEST_F(ParquetReaderTest, MismatchedSchemaFilterOnlyColumnCollision) auto const result = cudf::io::read_parquet(opts); CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); } + +TEST_F(ParquetReaderTest, RowIndexColumn) +{ + auto constexpr num_rows = 3; + auto col0 = cudf::test::fixed_width_column_wrapper( + cuda::counting_iterator{0}, cuda::counting_iterator{num_rows}); + auto col1 = cudf::test::fixed_width_column_wrapper( + cuda::counting_iterator{0}, cuda::counting_iterator{num_rows}); + auto table = cudf::table_view{{col0, col1}}; + + auto filepath = temp_env->get_temp_filepath("RowIndexColumn.parquet"); + cudf::io::write_parquet( + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, table).build()); + + auto const test_row_index_column = [&](auto num_sources) { + auto const sources = std::vector(num_sources, filepath); + auto read_opts = cudf::io::parquet_reader_options::builder(cudf::io::source_info(sources)) + .prepend_row_index_column(true) + .build(); + auto const read = cudf::io::read_parquet(read_opts); + EXPECT_EQ(read.tbl->num_columns(), table.num_columns() + 1); + EXPECT_EQ(read.metadata.schema_info.front().name, "row_idx"); + + // The row index restarts at zero for each source + auto const row_index = cudf::detail::make_counting_transform_iterator( + 0, [](cudf::size_type i) -> size_t { return i % num_rows; }); + auto const expected = + cudf::test::fixed_width_column_wrapper(row_index, row_index + num_sources * num_rows); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(read.tbl->view().column(0), expected); + + // The data columns are the single source file tiled `num_sources` times + auto const tiled = cudf::tile(table, num_sources); + CUDF_TEST_EXPECT_TABLES_EQUAL(read.tbl->select({1, 2}), tiled->view()); + + // Use filter `col0 < 2` so only rows {0, 1} from each source are read + auto scalar = cudf::numeric_scalar(2); + auto literal = cudf::ast::literal(scalar); + auto col_ref = cudf::ast::column_reference(0); + auto filter = cudf::ast::operation(cudf::ast::ast_operator::LESS, col_ref, literal); + + read_opts.set_filter(filter); + auto const read_filtered = cudf::io::read_parquet(read_opts); + auto const filtered_row_index = cudf::detail::make_counting_transform_iterator( + 0, [](cudf::size_type i) -> size_t { return i % 2; }); + auto const filtered_expected = cudf::test::fixed_width_column_wrapper( + filtered_row_index, filtered_row_index + num_sources * 2); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(read_filtered.tbl->view().column(0), filtered_expected); + }; + + test_row_index_column(1); + test_row_index_column(5); + + // With both source and row index columns enabled, the column order is: src_idx, row_idx, data + { + auto constexpr num_sources = 3; + auto const sources = std::vector(num_sources, filepath); + auto const read_opts = cudf::io::parquet_reader_options::builder(cudf::io::source_info(sources)) + .prepend_source_index_column(true) + .prepend_row_index_column(true) + .build(); + auto const read = cudf::io::read_parquet(read_opts); + EXPECT_EQ(read.tbl->num_columns(), table.num_columns() + 2); + EXPECT_EQ(read.metadata.schema_info[0].name, "src_idx"); + EXPECT_EQ(read.metadata.schema_info[1].name, "row_idx"); + + auto const src_index = cudf::detail::make_counting_transform_iterator( + 0, [](cudf::size_type i) { return i / num_rows; }); + auto const expected_src_index = cudf::test::fixed_width_column_wrapper( + src_index, src_index + num_sources * num_rows); + auto const row_index = cudf::detail::make_counting_transform_iterator( + 0, [](cudf::size_type i) -> size_t { return i % num_rows; }); + auto const expected_row_index = + cudf::test::fixed_width_column_wrapper(row_index, row_index + num_sources * num_rows); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(read.tbl->view().column(0), expected_src_index); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(read.tbl->view().column(1), expected_row_index); + } +} + +TEST_F(ParquetReaderTest, RowIndexSelectedRead) +{ + auto write_parquet = [](auto const& filepath, auto const& col) { + auto constexpr rows_per_row_group = 2; + auto const write_opts = cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, + cudf::table_view{{col}}) + .row_group_size_rows(rows_per_row_group) + .max_page_size_rows(rows_per_row_group) + .max_page_fragment_size(rows_per_row_group) + .build(); + cudf::io::write_parquet(write_opts); + }; + + auto col0 = cudf::test::fixed_width_column_wrapper{0, 1, 2, 3}; + auto filepath0 = temp_env->get_temp_filepath("RowIndexSelectedRead0.parquet"); + write_parquet(filepath0, col0); + + auto col1 = cudf::test::fixed_width_column_wrapper{10, 11, 12, 13}; + auto filepath1 = temp_env->get_temp_filepath("RowIndexSelectedRead1.parquet"); + write_parquet(filepath1, col1); + + auto source = cudf::io::source_info{std::vector{filepath0, filepath1}}; + + // The row index column contains file-local row indices + auto expected_values = cudf::test::fixed_width_column_wrapper{2, 3, 10, 11}; + auto expected_source = cudf::test::fixed_width_column_wrapper{0, 0, 1, 1}; + auto expected_row_index = cudf::test::fixed_width_column_wrapper{2, 3, 0, 1}; + auto expected = cudf::table_view{{expected_source, expected_row_index, expected_values}}; + + // Test with row bounds + { + auto const read_opts = cudf::io::parquet_reader_options::builder(source) + .skip_rows(2) + .num_rows(4) + .prepend_source_index_column(true) + .prepend_row_index_column(true) + .build(); + CUDF_TEST_EXPECT_TABLES_EQUAL(cudf::io::read_parquet(read_opts).tbl->view(), expected); + } + + // Test with row group selection + { + auto const read_opts = cudf::io::parquet_reader_options::builder(source) + .row_groups({{1}, {0}}) + .prepend_source_index_column(true) + .prepend_row_index_column(true) + .build(); + CUDF_TEST_EXPECT_TABLES_EQUAL(cudf::io::read_parquet(read_opts).tbl->view(), expected); + } + + // Test with out-of-order row group selection within a single source + { + auto const read_opts = cudf::io::parquet_reader_options::builder( + cudf::io::source_info{std::vector{filepath0}}) + .row_groups({{1, 0}}) + .prepend_row_index_column(true) + .build(); + auto const expected_ooo_values = cudf::test::fixed_width_column_wrapper{2, 3, 0, 1}; + auto const expected_ooo_row_index = cudf::test::fixed_width_column_wrapper{2, 3, 0, 1}; + auto const expected_ooo = cudf::table_view{{expected_ooo_row_index, expected_ooo_values}}; + CUDF_TEST_EXPECT_TABLES_EQUAL(cudf::io::read_parquet(read_opts).tbl->view(), expected_ooo); + } + + // Test with a filter: `col >= 2` prunes leading rows of source 0 (values {0, 1, 2, 3}) while + // keeping all of source 1 (values {10, 11, 12, 13}). The surviving rows of source 0 have + // file-local indices {2, 3}, verifying the row index reflects the original file position rather + // than the compacted output position. + { + auto scalar = cudf::numeric_scalar(2); + auto literal = cudf::ast::literal(scalar); + auto col_ref = cudf::ast::column_reference(0); + auto filter = cudf::ast::operation(cudf::ast::ast_operator::GREATER_EQUAL, col_ref, literal); + + auto const read_opts = cudf::io::parquet_reader_options::builder(source) + .filter(filter) + .prepend_source_index_column(true) + .prepend_row_index_column(true) + .build(); + auto const expected_filtered_values = + cudf::test::fixed_width_column_wrapper{2, 3, 10, 11, 12, 13}; + auto const expected_filtered_source = + cudf::test::fixed_width_column_wrapper{0, 0, 1, 1, 1, 1}; + auto const expected_filtered_row_index = + cudf::test::fixed_width_column_wrapper{2, 3, 0, 1, 2, 3}; + auto const expected_filtered = cudf::table_view{ + {expected_filtered_source, expected_filtered_row_index, expected_filtered_values}}; + CUDF_TEST_EXPECT_TABLES_EQUAL(cudf::io::read_parquet(read_opts).tbl->view(), expected_filtered); + } +} From f37d8e78265f63bf48e5fabaa6de73efd2c41a33 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Thu, 2 Jul 2026 19:14:22 +0000 Subject: [PATCH 2/7] Minor --- cpp/src/io/parquet/reader_impl.cpp | 5 +- cpp/src/io/parquet/reader_impl.hpp | 15 ++- cpp/src/io/parquet/reader_impl_preprocess.cu | 111 +++++++------------ 3 files changed, 53 insertions(+), 78 deletions(-) diff --git a/cpp/src/io/parquet/reader_impl.cpp b/cpp/src/io/parquet/reader_impl.cpp index ee123667f56c..e0644aa19e78 100644 --- a/cpp/src/io/parquet/reader_impl.cpp +++ b/cpp/src/io/parquet/reader_impl.cpp @@ -897,12 +897,13 @@ table_with_metadata reader_impl::finalize_output(read_mode mode, // Prepend the source and row index columns if requested { if (_options.prepend_row_index_column) { - prepend_row_index_column(read_info, out_columns); + 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_idx", .is_nullable = false}); } if (_options.prepend_source_index_column) { - prepend_source_index_column(out_metadata.num_rows_per_source, out_columns); + 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 = "src_idx", .is_nullable = false}); } diff --git a/cpp/src/io/parquet/reader_impl.hpp b/cpp/src/io/parquet/reader_impl.hpp index ef16c3487eef..81ae58ab05da 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -421,26 +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 num_rows_per_source, - std::vector>& out_columns); + [[nodiscard]] std::unique_ptr synthesize_source_index_column( + std::span num_rows_per_source); /** - * @brief Construct and prepend the file-local row index column to the output columns + * @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 - * @param out_columns Current output columns + * @return Synthesized row index column */ - void prepend_row_index_column(row_range const& read_info, - std::vector>& out_columns); + [[nodiscard]] std::unique_ptr synthesize_row_index_column(row_range const& read_info); /** * @brief Computes the names of columns to be read from the file, if specified. diff --git a/cpp/src/io/parquet/reader_impl_preprocess.cu b/cpp/src/io/parquet/reader_impl_preprocess.cu index 721f1c30f9fe..ba3eb995784d 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess.cu @@ -1131,68 +1131,55 @@ cudf::detail::host_vector reader_impl::calculate_page_string_offsets() namespace { /** - * @brief Maps each global row index in the current row range to the source-local row index + * @brief Maps each (global) row index in the row range to the source-local row index */ struct map_global_to_local_row_index { - std::size_t const* rg_global_row_offsets; ///< Global row offsets for each selected row group - std::size_t const* rg_local_row_offsets; ///< File-local row offsets for each selected row group + std::size_t const* global_row_offsets; ///< Global row offsets for each row group + std::size_t const* local_row_offsets; ///< Sources-local row offsets for each row group std::size_t num_row_groups; - __device__ std::size_t operator()(std::size_t global_row_index) const + __device__ std::size_t operator()(std::size_t row_idx) const noexcept { - auto const rg_idx = thrust::upper_bound(thrust::seq, - rg_global_row_offsets, - rg_global_row_offsets + num_row_groups, - global_row_index) - - rg_global_row_offsets - 1; - return global_row_index + rg_local_row_offsets[rg_idx]; + auto const row_group_idx = + thrust::upper_bound( + thrust::seq, global_row_offsets, global_row_offsets + num_row_groups, row_idx) - + global_row_offsets - 1; // Subtract 1 to get the index of the selected row group + return row_idx + local_row_offsets[row_group_idx]; } }; } // namespace -void reader_impl::prepend_row_index_column(row_range const& read_info, - std::vector>& out_columns) +std::unique_ptr reader_impl::synthesize_row_index_column(row_range const& read_info) { - if (not _options.prepend_row_index_column) { return; } - - using column_type = size_t; - auto constexpr dtype = cudf::data_type{cudf::type_to_id()}; - - auto const prepend_column = [&](auto&& column) { - out_columns.emplace(out_columns.begin(), std::move(column)); - }; - - // Empty column + using column_type = size_t; if (read_info.num_rows == 0) { - prepend_column(cudf::make_empty_column(dtype)); - return; + return cudf::make_empty_column(cudf::data_type{cudf::type_to_id()}); } - // Collect the global and file-local row offsets for each currently selected row group (in - // selection order). - auto const& row_groups = _file_itm_data.row_groups; - auto host_rg_global_offsets = - cudf::detail::make_empty_pinned_vector(row_groups.size(), _stream); - auto host_rg_local_offsets = - cudf::detail::make_empty_pinned_vector(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 - rg.start_row); - } + // Allocate column data vector + auto col_data = rmm::device_uvector(read_info.num_rows, _stream, _mr); // Map global row indices in the current row-range to corresponding source-local row indices - rmm::device_uvector col_data(read_info.num_rows, _stream, _mr); { - auto temp_mr = cudf::get_current_device_resource_ref(); + // Collect global and file-local row offsets 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(row_groups.size(), _stream); + auto host_rg_local_offsets = + cudf::detail::make_empty_pinned_vector(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 - rg.start_row); + } - // Copy row offsets to device - auto const rg_global_offsets = - cudf::detail::make_device_uvector_async(host_rg_global_offsets, _stream, temp_mr); - auto const rg_local_offsets = - cudf::detail::make_device_uvector_async(host_rg_local_offsets, _stream, temp_mr); + // 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 + // 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(read_info.skip_rows), col_data.begin(), @@ -1203,43 +1190,33 @@ void reader_impl::prepend_row_index_column(row_range const& read_info, _stream.synchronize(); } - prepend_column(std::make_unique(std::move(col_data), rmm::device_buffer{}, 0)); + return std::make_unique(std::move(col_data), rmm::device_buffer{}, 0); } -void reader_impl::prepend_source_index_column(std::span num_rows_per_source, - std::vector>& out_columns) +std::unique_ptr reader_impl::synthesize_source_index_column( + std::span num_rows_per_source) { - if (not _options.prepend_source_index_column) { return; } - - using column_type = cudf::size_type; - auto constexpr dtype = cudf::data_type{cudf::type_to_id()}; + 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()}); } // Single source - auto const num_sources = num_rows_per_source.size(); if (num_sources == 1) { auto const scalar = cudf::numeric_scalar(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(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(num_sources + 1, _stream); @@ -1247,15 +1224,13 @@ void reader_impl::prepend_source_index_column(std::span num_r 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); } - prepend_column(std::make_unique(std::move(col_data), rmm::device_buffer{}, 0)); + return std::make_unique(std::move(col_data), rmm::device_buffer{}, 0); } } // namespace cudf::io::parquet::detail From 178897572295214e0246eb04ebdf38217c7a71e6 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Tue, 7 Jul 2026 00:30:02 +0000 Subject: [PATCH 3/7] style --- cpp/tests/io/parquet_chunked_reader_test.cu | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/cpp/tests/io/parquet_chunked_reader_test.cu b/cpp/tests/io/parquet_chunked_reader_test.cu index 9edd60d471a3..0b4910ef4eeb 100644 --- a/cpp/tests/io/parquet_chunked_reader_test.cu +++ b/cpp/tests/io/parquet_chunked_reader_test.cu @@ -2113,13 +2113,11 @@ TEST_F(ParquetChunkedReaderTest, TestRowIndexColumnMultipleSources) } } - auto const expected_src_index = - cudf::test::fixed_width_column_wrapper(src_index_data.begin(), - src_index_data.end()); + auto const expected_src_index = cudf::test::fixed_width_column_wrapper( + src_index_data.begin(), src_index_data.end()); auto const expected_row_index = cudf::test::fixed_width_column_wrapper(row_index_data.begin(), row_index_data.end()); - auto const expected_values = - int64s_col(values_data.begin(), values_data.end()); + auto const expected_values = int64s_col(values_data.begin(), values_data.end()); CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view().column(0), expected_src_index); CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view().column(1), expected_row_index); CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view().column(2), expected_values); From cf68ff76ed1c7c1953f9165d895b1c38c797fcec Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Tue, 7 Jul 2026 01:26:42 +0000 Subject: [PATCH 4/7] missing stream sync --- cpp/src/io/parquet/reader_impl_preprocess.cu | 1 + 1 file changed, 1 insertion(+) diff --git a/cpp/src/io/parquet/reader_impl_preprocess.cu b/cpp/src/io/parquet/reader_impl_preprocess.cu index d837fadd058f..6a7e6ed2414f 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess.cu @@ -1228,6 +1228,7 @@ std::unique_ptr reader_impl::synthesize_source_index_column( 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(); } return std::make_unique(std::move(col_data), rmm::device_buffer{}, 0); From 15cc4fb73144183d8ef244150f0f71182d901a1c Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Tue, 7 Jul 2026 21:26:26 +0000 Subject: [PATCH 5/7] Minor improvements --- cpp/src/io/parquet/reader_impl.cpp | 7 +- cpp/src/io/parquet/reader_impl_preprocess.cu | 1 + cpp/tests/io/parquet_reader_test.cpp | 73 +++++++++++--------- 3 files changed, 46 insertions(+), 35 deletions(-) diff --git a/cpp/src/io/parquet/reader_impl.cpp b/cpp/src/io/parquet/reader_impl.cpp index e0644aa19e78..1155628159bf 100644 --- a/cpp/src/io/parquet/reader_impl.cpp +++ b/cpp/src/io/parquet/reader_impl.cpp @@ -899,13 +899,14 @@ table_with_metadata reader_impl::finalize_output(read_mode mode, 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_idx", .is_nullable = false}); + 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 = "src_idx", .is_nullable = false}); + out_metadata.schema_info.emplace( + out_metadata.schema_info.begin(), + column_name_info{.name = "source_index", .is_nullable = false}); } } diff --git a/cpp/src/io/parquet/reader_impl_preprocess.cu b/cpp/src/io/parquet/reader_impl_preprocess.cu index 6a7e6ed2414f..8b0679855f69 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess.cu @@ -1153,6 +1153,7 @@ struct map_global_to_local_row_index { std::unique_ptr reader_impl::synthesize_row_index_column(row_range const& read_info) { using column_type = size_t; + if (read_info.num_rows == 0) { return cudf::make_empty_column(cudf::data_type{cudf::type_to_id()}); } diff --git a/cpp/tests/io/parquet_reader_test.cpp b/cpp/tests/io/parquet_reader_test.cpp index 38462e08dc00..90ec77b33d7b 100644 --- a/cpp/tests/io/parquet_reader_test.cpp +++ b/cpp/tests/io/parquet_reader_test.cpp @@ -4802,7 +4802,7 @@ TEST_F(ParquetReaderTest, SourceIndexColumn) .build(); auto const read = cudf::io::read_parquet(read_opts); EXPECT_EQ(read.tbl->num_columns(), table.num_columns() + 1); - EXPECT_EQ(read.metadata.schema_info.front().name, "src_idx"); + EXPECT_EQ(read.metadata.schema_info.front().name, "source_index"); auto const src_index = cudf::detail::make_counting_transform_iterator( 0, [](cudf::size_type i) { return i / num_rows; }); @@ -4998,7 +4998,7 @@ TEST_F(ParquetReaderTest, RowIndexColumn) .build(); auto const read = cudf::io::read_parquet(read_opts); EXPECT_EQ(read.tbl->num_columns(), table.num_columns() + 1); - EXPECT_EQ(read.metadata.schema_info.front().name, "row_idx"); + EXPECT_EQ(read.metadata.schema_info.front().name, "row_index"); // The row index restarts at zero for each source auto const row_index = cudf::detail::make_counting_transform_iterator( @@ -5028,31 +5028,42 @@ TEST_F(ParquetReaderTest, RowIndexColumn) test_row_index_column(1); test_row_index_column(5); +} - // With both source and row index columns enabled, the column order is: src_idx, row_idx, data - { - auto constexpr num_sources = 3; - auto const sources = std::vector(num_sources, filepath); - auto const read_opts = cudf::io::parquet_reader_options::builder(cudf::io::source_info(sources)) - .prepend_source_index_column(true) - .prepend_row_index_column(true) - .build(); - auto const read = cudf::io::read_parquet(read_opts); - EXPECT_EQ(read.tbl->num_columns(), table.num_columns() + 2); - EXPECT_EQ(read.metadata.schema_info[0].name, "src_idx"); - EXPECT_EQ(read.metadata.schema_info[1].name, "row_idx"); +TEST_F(ParquetReaderTest, SourceAndRowIndexColumns) +{ + auto constexpr num_rows = 3; + auto constexpr num_sources = 3; + auto col0 = cudf::test::fixed_width_column_wrapper( + cuda::counting_iterator{0}, cuda::counting_iterator{num_rows}); + auto col1 = cudf::test::fixed_width_column_wrapper( + cuda::counting_iterator{0}, cuda::counting_iterator{num_rows}); + auto table = cudf::table_view{{col0, col1}}; - auto const src_index = cudf::detail::make_counting_transform_iterator( - 0, [](cudf::size_type i) { return i / num_rows; }); - auto const expected_src_index = cudf::test::fixed_width_column_wrapper( - src_index, src_index + num_sources * num_rows); - auto const row_index = cudf::detail::make_counting_transform_iterator( - 0, [](cudf::size_type i) -> size_t { return i % num_rows; }); - auto const expected_row_index = - cudf::test::fixed_width_column_wrapper(row_index, row_index + num_sources * num_rows); - CUDF_TEST_EXPECT_COLUMNS_EQUAL(read.tbl->view().column(0), expected_src_index); - CUDF_TEST_EXPECT_COLUMNS_EQUAL(read.tbl->view().column(1), expected_row_index); - } + auto filepath = temp_env->get_temp_filepath("RowIndexAndSourceIndexColumns.parquet"); + cudf::io::write_parquet( + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, table).build()); + + auto const sources = std::vector(num_sources, filepath); + auto const read_opts = cudf::io::parquet_reader_options::builder(cudf::io::source_info(sources)) + .prepend_source_index_column(true) + .prepend_row_index_column(true) + .build(); + auto const read = cudf::io::read_parquet(read_opts); + EXPECT_EQ(read.tbl->num_columns(), table.num_columns() + 2); + EXPECT_EQ(read.metadata.schema_info[0].name, "source_index"); + EXPECT_EQ(read.metadata.schema_info[1].name, "row_index"); + + auto const src_index = cudf::detail::make_counting_transform_iterator( + 0, [](cudf::size_type i) { return i / num_rows; }); + auto const expected_src_index = cudf::test::fixed_width_column_wrapper( + src_index, src_index + num_sources * num_rows); + auto const row_index = cudf::detail::make_counting_transform_iterator( + 0, [](cudf::size_type i) -> size_t { return i % num_rows; }); + auto const expected_row_index = + cudf::test::fixed_width_column_wrapper(row_index, row_index + num_sources * num_rows); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(read.tbl->view().column(0), expected_src_index); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(read.tbl->view().column(1), expected_row_index); } TEST_F(ParquetReaderTest, RowIndexSelectedRead) @@ -5112,16 +5123,14 @@ TEST_F(ParquetReaderTest, RowIndexSelectedRead) .row_groups({{1, 0}}) .prepend_row_index_column(true) .build(); - auto const expected_ooo_values = cudf::test::fixed_width_column_wrapper{2, 3, 0, 1}; - auto const expected_ooo_row_index = cudf::test::fixed_width_column_wrapper{2, 3, 0, 1}; - auto const expected_ooo = cudf::table_view{{expected_ooo_row_index, expected_ooo_values}}; - CUDF_TEST_EXPECT_TABLES_EQUAL(cudf::io::read_parquet(read_opts).tbl->view(), expected_ooo); + auto const expected_values = cudf::test::fixed_width_column_wrapper{2, 3, 0, 1}; + auto const expected_row_index = cudf::test::fixed_width_column_wrapper{2, 3, 0, 1}; + auto const expected = cudf::table_view{{expected_row_index, expected_values}}; + CUDF_TEST_EXPECT_TABLES_EQUAL(cudf::io::read_parquet(read_opts).tbl->view(), expected); } // Test with a filter: `col >= 2` prunes leading rows of source 0 (values {0, 1, 2, 3}) while - // keeping all of source 1 (values {10, 11, 12, 13}). The surviving rows of source 0 have - // file-local indices {2, 3}, verifying the row index reflects the original file position rather - // than the compacted output position. + // keeping all of source 1 (values {10, 11, 12, 13}). { auto scalar = cudf::numeric_scalar(2); auto literal = cudf::ast::literal(scalar); From 9f6acb5fc0957f7411000c70f28ab8ab7f1811c2 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Tue, 7 Jul 2026 22:01:24 +0000 Subject: [PATCH 6/7] Minor improvement --- cpp/src/io/parquet/reader_impl_preprocess.cu | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/cpp/src/io/parquet/reader_impl_preprocess.cu b/cpp/src/io/parquet/reader_impl_preprocess.cu index 8b0679855f69..69635fac407d 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess.cu @@ -1131,11 +1131,11 @@ cudf::detail::host_vector reader_impl::calculate_page_string_offsets() namespace { /** - * @brief Maps each (global) row index in the row range to the source-local row index + * @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; ///< Sources-local 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 @@ -1144,7 +1144,7 @@ struct map_global_to_local_row_index { thrust::upper_bound( thrust::seq, global_row_offsets, global_row_offsets + num_row_groups, row_idx) - global_row_offsets - 1; // Subtract 1 to get the index of the selected row group - return row_idx + local_row_offsets[row_group_idx]; + return row_idx - global_row_offsets[row_group_idx] + local_row_offsets[row_group_idx]; } }; @@ -1163,7 +1163,7 @@ std::unique_ptr reader_impl::synthesize_row_index_column(row_range const // Map global row indices in the current row-range to corresponding source-local row indices { - // Collect global and file-local row offsets for each selected row group + // 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(row_groups.size(), _stream); @@ -1171,7 +1171,7 @@ std::unique_ptr reader_impl::synthesize_row_index_column(row_range const cudf::detail::make_empty_pinned_vector(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 - rg.start_row); + host_rg_local_offsets.push_back(rg.source_start_row); } // Copy to device From 2574b9d1a2a1d2f2b5d6ca0f644cb4263553e61c Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Tue, 7 Jul 2026 22:03:27 +0000 Subject: [PATCH 7/7] Use `cuda::std::distance` --- cpp/src/io/parquet/reader_impl_preprocess.cu | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/cpp/src/io/parquet/reader_impl_preprocess.cu b/cpp/src/io/parquet/reader_impl_preprocess.cu index 69635fac407d..bc0540e5a919 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess.cu @@ -1141,9 +1141,11 @@ struct map_global_to_local_row_index { __device__ std::size_t operator()(std::size_t row_idx) const noexcept { auto const row_group_idx = - thrust::upper_bound( - thrust::seq, global_row_offsets, global_row_offsets + num_row_groups, row_idx) - - global_row_offsets - 1; // Subtract 1 to get the index of the selected row group + 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]; } };