diff --git a/cpp/include/cudf/io/parquet.hpp b/cpp/include/cudf/io/parquet.hpp index 6129ad0f78d4..05b2ed621909 100644 --- a/cpp/include/cudf/io/parquet.hpp +++ b/cpp/include/cudf/io/parquet.hpp @@ -109,6 +109,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; @@ -311,6 +313,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 * @@ -562,6 +578,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; } }; /** @@ -835,6 +858,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..1155628159bf 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,25 @@ 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) { + 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_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 = "source_index", .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..81ae58ab05da 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -421,13 +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 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 + * @return Synthesized row index column + */ + [[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. @@ -469,6 +481,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 37f11c1ddf46..b34c79fcd395 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 7f33cf249289..bc0540e5a919 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,40 +1128,98 @@ cudf::detail::host_vector reader_impl::calculate_page_string_offsets() return cudf::detail::make_pinned_vector(d_col_sizes, _stream); } -void reader_impl::prepend_source_index_column(std::span num_rows_per_source, - std::vector>& out_columns) +namespace { + +/** + * @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; ///< 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 + { + auto const row_group_idx = + 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]; + } +}; + +} // namespace + +std::unique_ptr reader_impl::synthesize_row_index_column(row_range const& read_info) { - if (not _options.prepend_source_index_column) { return; } + using column_type = size_t; + + if (read_info.num_rows == 0) { + return cudf::make_empty_column(cudf::data_type{cudf::type_to_id()}); + } + + // 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 + { + // 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); + 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); + } + + // 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 + 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(); + } - using column_type = cudf::size_type; - auto constexpr dtype = cudf::data_type{cudf::type_to_id()}; + return std::make_unique(std::move(col_data), rmm::device_buffer{}, 0); +} + +std::unique_ptr reader_impl::synthesize_source_index_column( + std::span num_rows_per_source) +{ + 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); @@ -1166,15 +1227,14 @@ 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); + _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); } } // namespace cudf::io::parquet::detail diff --git a/cpp/tests/io/parquet_chunked_reader_test.cu b/cpp/tests/io/parquet_chunked_reader_test.cu index d9eb03a867d6..0b4910ef4eeb 100644 --- a/cpp/tests/io/parquet_chunked_reader_test.cu +++ b/cpp/tests/io/parquet_chunked_reader_test.cu @@ -1983,6 +1983,147 @@ 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 4e1725623f77..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; }); @@ -4977,3 +4977,179 @@ 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_index"); + + // 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); +} + +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 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) +{ + 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_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}). + { + 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); + } +}