From 6165df40288e1d45315a358f260c9a3ac4243438 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Fri, 12 Jun 2026 23:10:10 +0000 Subject: [PATCH 01/15] Prepend source idx column to parquet --- cpp/include/cudf/io/parquet.hpp | 31 ++++++ cpp/include/cudf/io/types.hpp | 23 +--- .../parquet/expression_transform_helpers.cpp | 65 +++++++++++ .../parquet/expression_transform_helpers.hpp | 51 +++++++++ cpp/src/io/parquet/reader_impl.cpp | 60 ++++++----- cpp/src/io/parquet/reader_impl.hpp | 21 ++-- cpp/src/io/parquet/reader_impl_helpers.cpp | 2 +- cpp/src/io/parquet/reader_impl_preprocess.cu | 70 ++++++++++++ cpp/tests/io/parquet_chunked_reader_test.cu | 23 +++- cpp/tests/io/parquet_reader_test.cpp | 101 +++++++++++++++++- 10 files changed, 382 insertions(+), 65 deletions(-) diff --git a/cpp/include/cudf/io/parquet.hpp b/cpp/include/cudf/io/parquet.hpp index d17483bef6c8..dcd3fa28ea9a 100644 --- a/cpp/include/cudf/io/parquet.hpp +++ b/cpp/include/cudf/io/parquet.hpp @@ -106,6 +106,8 @@ class parquet_reader_options { // Whether column name matching is case sensitive. In case of multiple // case-insensitive matches, the first matched column is selected bool _case_sensitive_names = true; + // Whether to prepend a source file index column to the output + bool _prepend_source_index_column = false; std::optional> _reader_column_schema; @@ -298,6 +300,16 @@ class parquet_reader_options { */ [[nodiscard]] bool is_enabled_case_sensitive_names() const { return _case_sensitive_names; } + /** + * @brief Returns whether to prepend a source file index column to the output. + * + * @return `true` if a source file index column should be prepended + */ + [[nodiscard]] bool is_enabled_prepend_source_index_column() const + { + return _prepend_source_index_column; + } + /** * @brief Set a new source location * @@ -533,6 +545,13 @@ class parquet_reader_options { * @param val Boolean indicating whether to enable case-sensitive matching. */ void enable_case_sensitive_names(bool val) { _case_sensitive_names = val; } + + /** + * @brief Sets whether to prepend a source file index column to the output. + * + * @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; } }; /** @@ -796,6 +815,18 @@ class parquet_reader_options_builder { return *this; } + /** + * @brief Sets whether to prepend a source file index column to the output. + * + * @param val Boolean indicating whether to prepend a source file index column + * @return this for chaining + */ + parquet_reader_options_builder& prepend_source_index_column(bool val) + { + options._prepend_source_index_column = val; + return *this; + } + /** * @brief move parquet_reader_options member once it's built. */ diff --git a/cpp/include/cudf/io/types.hpp b/cpp/include/cudf/io/types.hpp index 4978960b7c30..0cdc6434840b 100644 --- a/cpp/include/cudf/io/types.hpp +++ b/cpp/include/cudf/io/types.hpp @@ -230,34 +230,13 @@ struct column_name_info { std::optional type_length; ///< Byte width of data (for fixed length data) std::vector children; ///< Child column names - /** - * @brief Construct a column name info with a name, optional nullabilty, and no children - * - * @param _name Column name - * @param _is_nullable True if column is nullable - * @param _is_binary True if column is binary data - */ - column_name_info(std::string _name, - std::optional _is_nullable = std::nullopt, - std::optional _is_binary = std::nullopt) - : name(std::move(_name)), is_nullable(_is_nullable), is_binary(_is_binary) - { - } - - column_name_info() = default; - /** * @brief Compares two column name info structs for equality * * @param rhs column name info struct to compare against * @return boolean indicating if this and rhs are equal */ - bool operator==(column_name_info const& rhs) const - { - return ((name == rhs.name) && (is_nullable == rhs.is_nullable) && - (is_binary == rhs.is_binary) && (type_length == rhs.type_length) && - (children == rhs.children)); - }; + bool operator==(column_name_info const& rhs) const = default; }; /** diff --git a/cpp/src/io/parquet/expression_transform_helpers.cpp b/cpp/src/io/parquet/expression_transform_helpers.cpp index f7931c98adb5..249786210e8f 100644 --- a/cpp/src/io/parquet/expression_transform_helpers.cpp +++ b/cpp/src/io/parquet/expression_transform_helpers.cpp @@ -355,4 +355,69 @@ std::optional>> collect_filtered_row_group_in return {filtered_row_group_indices}; } +offset_column_references::offset_column_references( + std::optional> expr, size_type offset) + : _offset{offset} +{ + if (!expr.has_value()) { return; } + + if (offset == 0) { + _converted_expr = expr; + return; + } + _converted_expr = expr.value().get().accept(*this); +} + +std::optional> +offset_column_references::get_converted_expr() const +{ + return _converted_expr; +} + +std::reference_wrapper offset_column_references::visit( + ast::literal const& expr) +{ + return expr; +} + +std::reference_wrapper offset_column_references::visit( + ast::column_reference const& expr) +{ + auto const& new_ref = _expr_tree.push( + ast::column_reference{expr.get_column_index() + _offset, expr.get_table_source()}); + return std::reference_wrapper(new_ref); +} + +std::reference_wrapper offset_column_references::visit( + ast::column_name_reference const&) +{ + CUDF_FAIL("Column name references are not supported in offset_column_references"); +} + +std::reference_wrapper offset_column_references::visit( + ast::operation const& expr) +{ + auto const operands = expr.get_operands(); + auto op = expr.get_operator(); + auto new_operands = visit_operands(operands); + auto const operator_arity = cudf::ast::detail::ast_operator_arity(op); + if (operator_arity == 2) { + _expr_tree.push(ast::operation{op, new_operands.front(), new_operands.back()}); + } else if (operator_arity == 1) { + _expr_tree.push(ast::operation{op, new_operands.front()}); + } + return std::reference_wrapper(_expr_tree.back()); +} + +std::vector> offset_column_references::visit_operands( + cudf::host_span const> operands) +{ + std::vector> transformed_operands; + for (auto const& operand : operands) { + auto const new_operand = operand.get().accept(*this); + transformed_operands.push_back(new_operand); + } + return transformed_operands; +} + } // namespace cudf::io::parquet::detail diff --git a/cpp/src/io/parquet/expression_transform_helpers.hpp b/cpp/src/io/parquet/expression_transform_helpers.hpp index 8af1722f215c..6d47e419c7aa 100644 --- a/cpp/src/io/parquet/expression_transform_helpers.hpp +++ b/cpp/src/io/parquet/expression_transform_helpers.hpp @@ -222,15 +222,18 @@ class named_to_reference_converter : public ast::detail::expression_transformer * @copydoc ast::detail::expression_transformer::visit(ast::literal const& ) */ std::reference_wrapper visit(ast::literal const& expr) override; + /** * @copydoc ast::detail::expression_transformer::visit(ast::column_reference const& ) */ std::reference_wrapper visit(ast::column_reference const& expr) override; + /** * @copydoc ast::detail::expression_transformer::visit(ast::column_name_reference const& ) */ std::reference_wrapper visit( ast::column_name_reference const& expr) override; + /** * @copydoc ast::detail::expression_transformer::visit(ast::operation const& ) */ @@ -311,6 +314,54 @@ class equality_literals_collector : public ast::detail::expression_transformer { cudf::host_span _schema_tree; }; +/** + * @brief Offsets every column referencein an expression by the specified value + * + */ +class offset_column_references : public ast::detail::expression_transformer { + public: + offset_column_references(std::optional> expr, + size_type offset); + + /** + * @copydoc ast::detail::expression_transformer::visit(ast::literal const& ) + */ + std::reference_wrapper visit(ast::literal const& expr) override; + + /** + * @copydoc ast::detail::expression_transformer::visit(ast::column_reference const& ) + */ + std::reference_wrapper visit(ast::column_reference const& expr) override; + + /** + * @copydoc ast::detail::expression_transformer::visit(ast::column_name_reference const& ) + */ + std::reference_wrapper visit( + ast::column_name_reference const& expr) override; + + /** + * @copydoc ast::detail::expression_transformer::visit(ast::operation const& ) + */ + std::reference_wrapper visit(ast::operation const& expr) override; + + /** + * @brief Returns the offset AST expression + * + * @return AST operation expression, or std::nullopt if no expression was provided + */ + [[nodiscard]] std::optional> get_converted_expr() + const; + + private: + std::vector> visit_operands( + cudf::host_span const> operands); + + size_type _offset{0}; + std::optional> _converted_expr; + // Owns the transformed expressions and keeps references stable + ast::tree _expr_tree; +}; + /** * @brief Maps indices of (all or selected) columns to their names * diff --git a/cpp/src/io/parquet/reader_impl.cpp b/cpp/src/io/parquet/reader_impl.cpp index eb00fafa6896..ef6683e2b353 100644 --- a/cpp/src/io/parquet/reader_impl.cpp +++ b/cpp/src/io/parquet/reader_impl.cpp @@ -512,7 +512,8 @@ reader_impl::reader_impl(std::size_t chunk_read_limit, options.get_num_bytes(), options.get_row_groups(), options.is_enabled_use_jit_filter(), - options.is_enabled_case_sensitive_names()}, + options.is_enabled_case_sensitive_names(), + options.is_enabled_prepend_source_index_column()}, _sources{std::move(sources)}, _output_chunk_read_limit{chunk_read_limit}, _input_pass_read_limit{pass_read_limit} @@ -678,12 +679,9 @@ table_with_metadata reader_impl::read_chunk_internal(read_mode mode) // no work to do (this can happen on the first pass if we have no rows to read) if (!has_more_work()) { - // Check if number of rows per source should be included in output metadata. - if (include_output_num_rows_per_source()) { - // Empty dataframe case: Simply initialize to a list of zeros - out_metadata.num_rows_per_source = - std::vector(_file_itm_data.num_rows_per_source.size(), 0); - } + // Empty dataframe case: Simply initialize to a list of zeros + out_metadata.num_rows_per_source = + std::vector(_file_itm_data.num_rows_per_source.size(), 0); // Finalize output return finalize_output(mode, out_metadata, out_columns); @@ -745,18 +743,13 @@ table_with_metadata reader_impl::read_chunk_internal(read_mode mode) out_columns = cudf::structs::detail::enforce_null_consistency(std::move(out_columns), _stream, _mr); - // Check if number of rows per source should be included in output metadata. - if (include_output_num_rows_per_source()) { - // For chunked reading, compute the output number of rows per source - if (mode == read_mode::CHUNKED_READ) { - out_metadata.num_rows_per_source = - calculate_output_num_rows_per_source(read_info.skip_rows, read_info.num_rows); - } - // Simply move the number of rows per file if reading all at once - else { - // Move is okay here as we are reading in one go. - out_metadata.num_rows_per_source = std::move(_file_itm_data.num_rows_per_source); - } + // Compute the output number of rows per source + if (mode == read_mode::CHUNKED_READ) { + out_metadata.num_rows_per_source = + calculate_output_num_rows_per_source(read_info.skip_rows, read_info.num_rows); + } else { + // Move is okay here as we are reading in one go. + out_metadata.num_rows_per_source = std::move(_file_itm_data.num_rows_per_source); } // Add empty columns if needed. Filter output columns based on filter. @@ -893,8 +886,21 @@ 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}); + } + + // 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 final_filter = + offset_column_references(_expr_conv.get_converted_expr(), num_prepended_cols); + auto const final_filter_expr = final_filter.get_converted_expr(); + // check if the output filter AST expression (= _expr_conv.get_converted_expr()) exists - if (_expr_conv.get_converted_expr().has_value()) { + if (final_filter_expr.has_value()) { auto read_table = std::make_unique(std::move(out_columns)); auto counting_it = cuda::counting_iterator{0}; auto const output_count = read_table->num_columns() - _num_filter_only_columns; @@ -902,11 +908,14 @@ table_with_metadata reader_impl::finalize_output(read_mode mode, if (_num_filter_only_columns > 0) { out_metadata.schema_info.resize(output_count); } + // Clear the number of rows per source as it is not valid after filtering + out_metadata.num_rows_per_source.clear(); + bool use_jit = cudf::get_context().use_jit() || _options.use_jit_filter; if (!use_jit) { - auto predicate = cudf::detail::compute_column( - *read_table, _expr_conv.get_converted_expr().value().get(), _stream, _mr); + auto predicate = + cudf::detail::compute_column(*read_table, final_filter_expr.value().get(), _stream, _mr); CUDF_EXPECTS(predicate->view().type().id() == type_id::BOOL8, "Predicate filter should return a boolean"); // Exclude columns present in filter only in output @@ -914,11 +923,8 @@ table_with_metadata reader_impl::finalize_output(read_mode mode, only_output, *predicate, cudf::detail::mask_type::RETENTION, _stream, _mr); return {std::move(output_table), std::move(out_metadata)}; } else { - auto output_table = cudf::filter(read_table->view(), - _expr_conv.get_converted_expr().value().get(), - only_output, - _stream, - _mr); + auto output_table = cudf::filter( + read_table->view(), final_filter_expr.value().get(), only_output, _stream, _mr); return {std::move(output_table), std::move(out_metadata)}; } diff --git a/cpp/src/io/parquet/reader_impl.hpp b/cpp/src/io/parquet/reader_impl.hpp index 4413185e926e..9580c0590a2b 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -395,16 +395,6 @@ class reader_impl { return _file_itm_data._output_chunk_count == 0; } - /** - * @brief Check if number of rows per source should be included in output metadata. - * - * @return True if AST filter is not present - */ - [[nodiscard]] bool include_output_num_rows_per_source() const - { - return not _expr_conv.get_converted_expr().has_value(); - } - [[nodiscard]] cudf::detail::hostdevice_span subpass_page_mask_span() const { return _subpass_page_mask ? *_subpass_page_mask : cudf::detail::hostdevice_span{}; @@ -420,6 +410,15 @@ class reader_impl { [[nodiscard]] std::vector calculate_output_num_rows_per_source(size_t chunk_start_row, size_t chunk_num_rows); + /** + * @brief Construct and prepend the source index column to the output columns + * + * @param num_rows_per_source Number of rows per parquet source + * @param out_columns Current output columns + */ + void prepend_source_index_column(std::span num_rows_per_source, + std::vector>& out_columns); + /** * @brief Computes the names of columns to be read from the file, if specified. * @@ -458,6 +457,8 @@ class reader_impl { bool use_jit_filter = false; // Whether to use case-sensitive matching for column names bool case_sensitive_names = true; + // Whether to prepend the source file index column to the output + bool prepend_source_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 7208e2e32091..31926854f95a 100644 --- a/cpp/src/io/parquet/reader_impl_helpers.cpp +++ b/cpp/src/io/parquet/reader_impl_helpers.cpp @@ -1970,7 +1970,7 @@ aggregate_reader_metadata::select_columns( std::transform(index_names.cbegin(), index_names.cend(), std::back_inserter(selected_columns), - [](std::string const& name) { return column_name_info(name); }); + [](std::string const& name) { return column_name_info{.name = name}; }); } // Merge the vector use_names into a set of hierarchical column_name_info objects /* This is because if we have columns like this: diff --git a/cpp/src/io/parquet/reader_impl_preprocess.cu b/cpp/src/io/parquet/reader_impl_preprocess.cu index 31f25f49d0e8..63b97549ee13 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess.cu @@ -9,6 +9,7 @@ #include "reader_impl_chunking_utils.cuh" #include "reader_impl_preprocess_utils.cuh" +#include #include #include #include @@ -1118,4 +1119,73 @@ 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) +{ + if (not _options.prepend_source_index_column) { return; } + + using column_type = cudf::size_type; + auto constexpr dtype = cudf::data_type{cudf::type_to_id()}; + + auto const num_rows = std::accumulate(num_rows_per_source.begin(), num_rows_per_source.end(), 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; + } + + // 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; + } + + auto col_data = + cudf::detail::make_zeroed_device_uvector_async(num_rows, _stream, _mr); + + // Multiple sources + { + auto temp_mr = cudf::get_current_device_resource_ref(); + + // Host per-source exclusive row offsets + auto host_row_offsets = + cudf::detail::make_empty_pinned_vector(num_sources, _stream); + host_row_offsets.resize(num_sources); + std::exclusive_scan(num_rows_per_source.begin(), + num_rows_per_source.end(), + host_row_offsets.begin(), + cudf::size_type{0}); + + // Copy row counts (scatter stencil) and offsets (scatter map) to device. + auto const row_counts = + cudf::detail::make_device_uvector_async(num_rows_per_source, _stream, temp_mr); + auto const row_offsets = + cudf::detail::make_device_uvector_async(host_row_offsets, _stream, temp_mr); + + // Scatter each source index to its first row's position, but only for sources that have rows. + thrust::scatter_if(rmm::exec_policy_nosync(_stream, temp_mr), + cuda::counting_iterator(0), + cuda::counting_iterator(num_sources), + row_offsets.begin(), + row_counts.begin(), + col_data.begin()); + // Inclusive scan with maximum to replace zeros with the source index each row belongs to. + thrust::inclusive_scan(rmm::exec_policy_nosync(_stream, temp_mr), + col_data.begin(), + col_data.end(), + col_data.begin(), + cuda::maximum()); + _stream.synchronize(); + } + + prepend_column(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 489b95126568..69603393d88a 100644 --- a/cpp/tests/io/parquet_chunked_reader_test.cu +++ b/cpp/tests/io/parquet_chunked_reader_test.cu @@ -1733,7 +1733,7 @@ TEST_F(ParquetChunkedReaderTest, TestNumRowsPerSource) EXPECT_EQ(num_rows_per_source[0], num_rows - rows_to_skip); } - // Filtered chunked-read from single data source + // Filtered chunked-read from single data source with prepended source index column { int64_t const max_value = int64_data[int64_data.size() / 2]; auto constexpr output_read_limit = 1'500; @@ -1746,6 +1746,7 @@ TEST_F(ParquetChunkedReaderTest, TestNumRowsPerSource) auto const options = cudf::io::parquet_reader_options_builder(cudf::io::source_info{filepath}) .filter(filter_expression) + .prepend_source_index_column(true) .build(); auto const reader = cudf::io::chunked_parquet_reader( output_read_limit, pass_read_limit, options, cudf::get_default_stream()); @@ -1764,7 +1765,12 @@ TEST_F(ParquetChunkedReaderTest, TestNumRowsPerSource) cudf::table_view expected_filtered({int64_col_filtered->view()}); - CUDF_TEST_EXPECT_TABLES_EQUAL(expected_filtered, result->view()); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected_filtered, result->select({1})); + // Expected source index column + auto const zeros = cuda::make_constant_iterator(0); + auto const expected_src_index = cudf::test::fixed_width_column_wrapper( + zeros, zeros + int64_data_filtered.size()); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view().column(0), expected_src_index); EXPECT_TRUE(num_rows_per_source.empty()); } } @@ -1822,7 +1828,7 @@ TEST_F(ParquetChunkedReaderTest, TestNumRowsPerSourceMultipleSources) return expected_counts; }; - // Chunked-read six data sources entirely + // Chunked-read six data sources entirely with prepended source index column { auto const nsources = 6; auto constexpr output_read_limit = 15'000; @@ -1830,7 +1836,9 @@ TEST_F(ParquetChunkedReaderTest, TestNumRowsPerSourceMultipleSources) std::vector const datasources(nsources, filepath); auto const options = - cudf::io::parquet_reader_options_builder(cudf::io::source_info{datasources}).build(); + cudf::io::parquet_reader_options_builder(cudf::io::source_info{datasources}) + .prepend_source_index_column(true) + .build(); auto const reader = cudf::io::chunked_parquet_reader( output_read_limit, pass_read_limit, options, cudf::get_default_stream()); @@ -1842,6 +1850,13 @@ TEST_F(ParquetChunkedReaderTest, TestNumRowsPerSourceMultipleSources) EXPECT_EQ(num_rows_per_source.size(), nsources); EXPECT_TRUE( std::equal(expected_counts.cbegin(), expected_counts.cend(), num_rows_per_source.cbegin())); + + // Expected source index column + 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); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view().column(0), expected_src_index); } // Chunked-read rows_to_read rows skipping rows_to_skip from eight data sources diff --git a/cpp/tests/io/parquet_reader_test.cpp b/cpp/tests/io/parquet_reader_test.cpp index 8ca9de9c10db..8a43b10396a1 100644 --- a/cpp/tests/io/parquet_reader_test.cpp +++ b/cpp/tests/io/parquet_reader_test.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -1495,9 +1496,10 @@ TEST_F(ParquetReaderTest, FilterWithColumnProjection) auto read_opts = cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}) .column_names({"col_Double"}) .case_sensitive_names(false) + .prepend_source_index_column(true) .filter(read_expr); auto result = cudf::io::read_parquet(read_opts); - CUDF_TEST_EXPECT_TABLES_EQUAL(*result.tbl, *expected); + CUDF_TEST_EXPECT_TABLES_EQUAL(*result.tbl->select({1}), *expected); // Repeat but select columns using indices instead of names read_opts = cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}) @@ -4655,3 +4657,100 @@ TEST_F(ParquetReaderTest, DecodeVariableWidthDecimalStats) EXPECT_THROW(ParquetStatsDecoder::decode({0xff, 0xff, 0xff, 0xfe, 0x0c}), cudf::logic_error); } + +TEST_F(ParquetReaderTest, SourceIndexColumn) +{ + 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("SourceIndexNoFilter.parquet"); + cudf::io::write_parquet( + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, table).build()); + + auto const test_source_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_source_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, "src_idx"); + + auto const src_index = cudf::detail::make_counting_transform_iterator( + 0, [](cudf::size_type i) { return i / num_rows; }); + auto const expected = cudf::test::fixed_width_column_wrapper( + src_index, src_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_src_index = + cudf::detail::make_counting_transform_iterator(0, [](cudf::size_type i) { return i / 2; }); + auto const filtered_expected = cudf::test::fixed_width_column_wrapper( + filtered_src_index, filtered_src_index + num_sources * 2); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(read_filtered.tbl->view().column(0), filtered_expected); + }; + + test_source_index_column(1); + test_source_index_column(5); +} + +TEST_F(ParquetReaderTest, SourceIndexSelectedRead) +{ + 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("SourceIndexSelectedRead0.parquet"); + write_parquet(filepath0, col0); + + auto col1 = cudf::test::fixed_width_column_wrapper{10, 11, 12, 13}; + auto filepath1 = temp_env->get_temp_filepath("SourceIndexSelectedRead1.parquet"); + write_parquet(filepath1, col1); + + auto source = cudf::io::source_info{std::vector{filepath0, filepath1}}; + + 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 = cudf::table_view{{expected_source, 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) + .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) + .build(); + CUDF_TEST_EXPECT_TABLES_EQUAL(cudf::io::read_parquet(read_opts).tbl->view(), expected); + } +} \ No newline at end of file From e6827d716223b370d9c04c373b02ef4f7f53e316 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Fri, 12 Jun 2026 23:11:12 +0000 Subject: [PATCH 02/15] Prepend source index column to hybrid scan --- .../parquet/experimental/hybrid_scan_impl.cpp | 74 +++++---- .../io/experimental/hybrid_scan_test.cpp | 142 ++++++++++++++++++ cpp/tests/io/parquet_reader_test.cpp | 2 +- 3 files changed, 184 insertions(+), 34 deletions(-) diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index 0f0fd7de6d96..6b399343639a 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -805,12 +805,13 @@ void hybrid_scan_reader_impl::reset_internal_state() _options.decimal_width = type_id::EMPTY; _options.num_rows = std::nullopt; _options.row_group_indices.clear(); - _options.use_jit_filter = false; - _options.case_sensitive_names = true; - _num_sources = 0; - _input_pass_read_limit = 0; - _output_chunk_read_limit = 0; - _strings_to_categorical = false; + _options.use_jit_filter = false; + _options.case_sensitive_names = true; + _options.prepend_source_index_column = false; + _num_sources = 0; + _input_pass_read_limit = 0; + _output_chunk_read_limit = 0; + _strings_to_categorical = false; _reader_column_schema.reset(); _expr_conv = named_to_reference_converter{}; _mr = cudf::get_current_device_resource_ref(); @@ -822,10 +823,11 @@ void hybrid_scan_reader_impl::initialize_column_selection_options( // Strings may be returned as either string or categorical columns _strings_to_categorical = options.is_enabled_convert_strings_to_categories(); - _options.timestamp_type = cudf::data_type{options.get_timestamp_type().id()}; - _options.decimal_width = options.get_decimal_width(); - _options.use_jit_filter = options.is_enabled_use_jit_filter(); - _options.case_sensitive_names = options.is_enabled_case_sensitive_names(); + _options.timestamp_type = cudf::data_type{options.get_timestamp_type().id()}; + _options.decimal_width = options.get_decimal_width(); + _options.use_jit_filter = options.is_enabled_use_jit_filter(); + _options.case_sensitive_names = options.is_enabled_case_sensitive_names(); + _options.prepend_source_index_column = options.is_enabled_prepend_source_index_column(); _use_pandas_metadata = options.is_enabled_use_pandas_metadata(); } @@ -914,12 +916,9 @@ table_with_metadata hybrid_scan_reader_impl::read_chunk_internal( // no work to do (this can happen on the first pass if we have no rows to read) if (!has_more_work()) { - // Check if number of rows per source should be included in output metadata. - if (include_output_num_rows_per_source()) { - // Empty dataframe case: Simply initialize to a list of zeros - out_metadata.num_rows_per_source = - std::vector(_file_itm_data.num_rows_per_source.size(), 0); - } + // Empty dataframe case: Simply initialize to a list of zeros + out_metadata.num_rows_per_source = + std::vector(_file_itm_data.num_rows_per_source.size(), 0); // Finalize output return finalize_output(read_columns_mode, out_metadata, out_columns, row_mask); @@ -981,18 +980,13 @@ table_with_metadata hybrid_scan_reader_impl::read_chunk_internal( out_columns = cudf::structs::detail::enforce_null_consistency(std::move(out_columns), _stream, _mr); - // Check if number of rows per source should be included in output metadata. - if (include_output_num_rows_per_source()) { - // For chunked reading, compute the output number of rows per source - if (mode == read_mode::CHUNKED_READ) { - out_metadata.num_rows_per_source = - calculate_output_num_rows_per_source(read_info.skip_rows, read_info.num_rows); - } - // Simply move the number of rows per file if reading all at once - else { - // Move is okay here as we are reading in one go. - out_metadata.num_rows_per_source = std::move(_file_itm_data.num_rows_per_source); - } + // Compute the output number of rows per source + if (mode == read_mode::CHUNKED_READ) { + out_metadata.num_rows_per_source = + calculate_output_num_rows_per_source(read_info.skip_rows, read_info.num_rows); + } else { + // Move is okay here as we are reading in one go. + out_metadata.num_rows_per_source = std::move(_file_itm_data.num_rows_per_source); } // Add empty columns if needed. Filter output columns based on filter. @@ -1035,6 +1029,18 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output( apply_decimal_width_cast(out_columns); + // Prepend the source index column to filter columns or all columns + if (_options.prepend_source_index_column and + read_columns_mode == read_columns_mode::FILTER_COLUMNS) { + prepend_source_index_column(out_metadata.num_rows_per_source, out_columns); + } + + // 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 final_filter = + parquet::detail::offset_column_references(_expr_conv.get_converted_expr(), num_prepended_cols); + auto const final_filter_expr = final_filter.get_converted_expr(); + // Create a table from the output columns. auto read_table = std::make_unique
(std::move(out_columns)); @@ -1050,15 +1056,17 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output( _row_mask_offset += read_table->num_rows(); _output_chunk_produced = true; + // Clear the number of rows per source as it is not valid after filtering + out_metadata.num_rows_per_source.clear(); + // For filter columns, apply the filter expression and update the input row mask if constexpr (std::is_same_v) { CUDF_EXPECTS(read_columns_mode == read_columns_mode::FILTER_COLUMNS, "Invalid read mode"); - auto final_row_mask = - cudf::detail::compute_column(*read_table, - _expr_conv.get_converted_expr().value().get(), - _stream, - cudf::get_current_device_resource_ref()); + auto final_row_mask = cudf::detail::compute_column(*read_table, + final_filter_expr.value().get(), + _stream, + cudf::get_current_device_resource_ref()); CUDF_EXPECTS(final_row_mask->view().type().id() == type_id::BOOL8, "Predicate filter should return a boolean"); diff --git a/cpp/tests/io/experimental/hybrid_scan_test.cpp b/cpp/tests/io/experimental/hybrid_scan_test.cpp index 05e7fba3d844..d7004f127592 100644 --- a/cpp/tests/io/experimental/hybrid_scan_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_test.cpp @@ -923,6 +923,148 @@ TEST_F(HybridScanTest, StructChildFilterColumn) std::invalid_argument); } +TEST_F(HybridScanTest, PrependSourceIndexColumn) +{ + using T = uint32_t; + using cudf::io::parquet::experimental::use_data_page_mask; + + // Single source parquet file with columns col0 (filter), col1 and col2 (payload) + auto constexpr num_concat = 1; + auto [written_table, parquet_buffer] = create_parquet_with_stats(); + + // Filtering AST - col0 < 100 + auto literal_value = cudf::numeric_scalar(100); + auto literal = cudf::ast::literal(literal_value); + auto col_ref_0 = cudf::ast::column_name_reference("col0"); + auto filter_expression = cudf::ast::operation(cudf::ast::ast_operator::LESS, col_ref_0, literal); + + auto const stream = cudf::get_default_stream(); + auto const mr = cudf::get_current_device_resource_ref(); + + // Expected table + auto const expected = [&]() { + auto const options = + cudf::io::parquet_reader_options::builder( + cudf::io::source_info(cudf::host_span(parquet_buffer.data(), parquet_buffer.size()))) + .filter(filter_expression) + .prepend_source_index_column(true) + .build(); + return cudf::io::read_parquet(options, stream, mr).tbl; + }(); + + // Hybrid scan reader options + auto const options = cudf::io::parquet_reader_options::builder() + .filter(filter_expression) + .prepend_source_index_column(true) + .build(); + + auto datasource = cudf::io::datasource::create(cudf::host_span( + reinterpret_cast(parquet_buffer.data()), parquet_buffer.size())); + auto const footer_buffer = cudf::io::parquet::fetch_footer_to_host(*datasource); + + auto reader = + std::make_unique(*footer_buffer, options); + + auto const all_row_groups = reader->all_row_groups(options); + + auto const current_row_groups = cudf::host_span(all_row_groups); + auto const mid = current_row_groups.size() / 2; + auto const row_group_passes = + std::array{current_row_groups.subspan(0, mid), + current_row_groups.subspan(mid, current_row_groups.size() - mid)}; + auto const table_views = [](auto const& tables) { + auto views = std::vector{}; + std::transform(tables.begin(), tables.end(), std::back_inserter(views), [](auto const& t) { + return t->view(); + }); + return views; + }; + + // Two-step chunked read + { + auto filter_tables = std::vector>{}; + auto payload_tables = std::vector>{}; + + auto const read_pass = [&](auto const row_groups) { + auto row_mask = reader->build_all_true_row_mask(row_groups, stream, mr); + auto row_mask_view = row_mask->mutable_view(); + + auto const filter_byte_ranges = reader->filter_column_chunks_byte_ranges(row_groups, options); + auto [filter_bufs, filter_data, filter_tasks] = + cudf::io::parquet::fetch_byte_ranges_to_device_async( + *datasource, filter_byte_ranges, stream, mr); + filter_tasks.get(); + reader->setup_chunking_for_filter_columns(1024, + 10240, + row_groups, + row_mask_view, + use_data_page_mask::NO, + filter_data, + options, + stream, + mr); + while (reader->has_next_table_chunk()) { + filter_tables.push_back(reader->materialize_filter_columns_chunk(row_mask_view).tbl); + } + + auto const payload_byte_ranges = + reader->payload_column_chunks_byte_ranges(row_groups, options); + auto [payload_bufs, payload_data, payload_tasks] = + cudf::io::parquet::fetch_byte_ranges_to_device_async( + *datasource, payload_byte_ranges, stream, mr); + payload_tasks.get(); + + reader->setup_chunking_for_payload_columns(1024, + 10240, + row_groups, + row_mask_view, + use_data_page_mask::NO, + payload_data, + options, + stream, + mr); + while (reader->has_next_table_chunk()) { + payload_tables.push_back(reader->materialize_payload_columns_chunk(row_mask_view).tbl); + } + }; + + for (auto const& pass : row_group_passes) { + read_pass(pass); + } + + auto const filter_table = cudf::concatenate(table_views(filter_tables), stream, mr); + auto const payload_table = cudf::concatenate(table_views(payload_tables), stream, mr); + + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->select({0, 1}), filter_table->view()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->select({2, 3}), payload_table->view()); + } + + // Single-step chunked read + { + auto all_tables = std::vector>{}; + + auto const read_pass = [&](cudf::host_span row_groups) { + auto const chunk_ranges = reader->all_column_chunks_byte_ranges(row_groups, options); + auto [bufs, col_data, tasks] = + cudf::io::parquet::fetch_byte_ranges_to_device_async(*datasource, chunk_ranges, stream, mr); + tasks.get(); + + reader->setup_chunking_for_all_columns( + 1024, 10240, row_groups, col_data, options, stream, mr); + while (reader->has_next_table_chunk()) { + all_tables.push_back(reader->materialize_all_columns_chunk().tbl); + } + }; + + for (auto const& pass : row_group_passes) { + read_pass(pass); + } + + auto const all_columns_table = cudf::concatenate(table_views(all_tables), stream, mr); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), all_columns_table->view()); + } +} + TEST_F(HybridScanTest, ChunkedReadRowMaskPerPass) { using T = uint32_t; diff --git a/cpp/tests/io/parquet_reader_test.cpp b/cpp/tests/io/parquet_reader_test.cpp index 8a43b10396a1..e60a03409358 100644 --- a/cpp/tests/io/parquet_reader_test.cpp +++ b/cpp/tests/io/parquet_reader_test.cpp @@ -1499,7 +1499,7 @@ TEST_F(ParquetReaderTest, FilterWithColumnProjection) .prepend_source_index_column(true) .filter(read_expr); auto result = cudf::io::read_parquet(read_opts); - CUDF_TEST_EXPECT_TABLES_EQUAL(*result.tbl->select({1}), *expected); + CUDF_TEST_EXPECT_TABLES_EQUAL(result.tbl->select({1}), expected->view()); // Repeat but select columns using indices instead of names read_opts = cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}) From 00b962178673fd70bd3f883a4fcdf1212f37d7be Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Fri, 12 Jun 2026 23:19:17 +0000 Subject: [PATCH 03/15] Style --- cpp/include/cudf/io/types.hpp | 2 +- cpp/tests/io/parquet_reader_test.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/include/cudf/io/types.hpp b/cpp/include/cudf/io/types.hpp index 0cdc6434840b..46dc25ed3f37 100644 --- a/cpp/include/cudf/io/types.hpp +++ b/cpp/include/cudf/io/types.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/io/parquet_reader_test.cpp b/cpp/tests/io/parquet_reader_test.cpp index e60a03409358..50a201faf2a5 100644 --- a/cpp/tests/io/parquet_reader_test.cpp +++ b/cpp/tests/io/parquet_reader_test.cpp @@ -4753,4 +4753,4 @@ TEST_F(ParquetReaderTest, SourceIndexSelectedRead) .build(); CUDF_TEST_EXPECT_TABLES_EQUAL(cudf::io::read_parquet(read_opts).tbl->view(), expected); } -} \ No newline at end of file +} From a0159e98be47625136048aadf232c6d712143bca Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Wed, 1 Jul 2026 12:39:59 -0700 Subject: [PATCH 04/15] Apply suggestion from @mhaseeb123 --- cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index 215a9f2679f3..dcd127f42cc8 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -1084,7 +1084,7 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output( apply_decimal_width_cast(out_columns); - // Prepend the source index column to filter columns or all columns + // Prepend the source index column to filter columns only if (_options.prepend_source_index_column and read_columns_mode == read_columns_mode::FILTER_COLUMNS) { prepend_source_index_column(out_metadata.num_rows_per_source, out_columns); From ceedee0149e413578e97d44cee27564cca1aeb5e Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Wed, 1 Jul 2026 20:54:22 +0000 Subject: [PATCH 05/15] Address comments --- .../parquet/experimental/hybrid_scan_impl.cpp | 16 +- cpp/src/io/parquet/reader_impl.cpp | 6 +- cpp/src/io/parquet/reader_impl.hpp | 12 ++ .../hybrid_scan_multifile_test.cpp | 75 +++++++++ .../io/experimental/hybrid_scan_test.cpp | 142 ------------------ cpp/tests/strings/contains_tests.cpp | 2 +- cpp/tests/strings/findall_tests.cpp | 2 +- cpp/tests/strings/replace_regex_tests.cpp | 2 +- 8 files changed, 100 insertions(+), 157 deletions(-) diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index dcd127f42cc8..9db3bf049ce5 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -856,6 +856,7 @@ void hybrid_scan_reader_impl::reset_internal_state() _pass_page_mask.clear(); _subpass_page_mask.reset(); _output_metadata.reset(); + _options.timestamp_type = cudf::data_type{}; _options.decimal_width = type_id::EMPTY; _options.num_rows = std::nullopt; @@ -863,10 +864,11 @@ void hybrid_scan_reader_impl::reset_internal_state() _options.use_jit_filter = false; _options.case_sensitive_names = true; _options.prepend_source_index_column = false; - _num_sources = 0; - _input_pass_read_limit = 0; - _output_chunk_read_limit = 0; - _strings_to_categorical = false; + + _num_sources = 0; + _input_pass_read_limit = 0; + _output_chunk_read_limit = 0; + _strings_to_categorical = false; _reader_column_schema.reset(); _expr_conv = named_to_reference_converter{}; _mr = cudf::get_current_device_resource_ref(); @@ -1090,10 +1092,8 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output( prepend_source_index_column(out_metadata.num_rows_per_source, out_columns); } - // 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 final_filter = - parquet::detail::offset_column_references(_expr_conv.get_converted_expr(), num_prepended_cols); + // Compute the final filter expression incorporating any column reference offsets in _expr_conv + auto const final_filter = compute_offset_filter(); auto const final_filter_expr = final_filter.get_converted_expr(); // Create a table from the output columns. diff --git a/cpp/src/io/parquet/reader_impl.cpp b/cpp/src/io/parquet/reader_impl.cpp index aef720839fa2..d9d247b6b325 100644 --- a/cpp/src/io/parquet/reader_impl.cpp +++ b/cpp/src/io/parquet/reader_impl.cpp @@ -893,10 +893,8 @@ table_with_metadata reader_impl::finalize_output(read_mode mode, 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 final_filter = - offset_column_references(_expr_conv.get_converted_expr(), num_prepended_cols); + // Compute the final filter expression incorporating any column reference offsets in _expr_conv + auto const final_filter = compute_offset_filter(); auto const final_filter_expr = final_filter.get_converted_expr(); // check if the output filter AST expression (= _expr_conv.get_converted_expr()) exists diff --git a/cpp/src/io/parquet/reader_impl.hpp b/cpp/src/io/parquet/reader_impl.hpp index 82bbe422cbfa..32b96ec56b2a 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -419,6 +419,18 @@ class reader_impl { void prepend_source_index_column(std::span num_rows_per_source, std::vector>& out_columns); + /** + * @brief Offset the column references in `_expr_conv` by the number of columns prepended to + * the output + * + * @return Offsetted expression converter + */ + [[nodiscard]] offset_column_references compute_offset_filter() const + { + auto const num_prepended_cols = static_cast(_options.prepend_source_index_column); + return offset_column_references(_expr_conv.get_converted_expr(), num_prepended_cols); + } + /** * @brief Computes the names of columns to be read from the file, if specified. * diff --git a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp index 65b04eedfdd5..3fcecb1c2fd8 100644 --- a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include #include #include @@ -158,6 +159,80 @@ TEST_F(HybridScanMultifileTest, MaterializeListsOfStrings) test_hybrid_scan_multifile({col0, *col1, *col2, *col3, *col4}, false); } +TEST_F(HybridScanMultifileTest, PrependSourceIndexColumn) +{ + using T = int32_t; + using cudf::io::parquet::experimental::use_data_page_mask; + + // Small single-column table with sequence values [0, 10) + auto constexpr num_rows = 10; + auto constexpr num_sources = 3; + auto values = cuda::counting_iterator{0}; + cudf::test::fixed_width_column_wrapper col0(values, values + num_rows); + auto const table = cudf::table_view{{col0}}; + + cudf::io::table_input_metadata expected_metadata(table); + expected_metadata.column_metadata[0].set_name("col0"); + + std::vector> parquet_buffers(num_sources); + for (auto& parquet_buffer : parquet_buffers) { + auto out_opts = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&parquet_buffer}, table) + .metadata(expected_metadata) + .build(); + cudf::io::write_parquet(out_opts); + } + + // Filtering AST - col0 % 2 == 0, removes odd rows (half the rows) from each source + auto two_scalar = cudf::numeric_scalar(2); + auto two_literal = cudf::ast::literal(two_scalar); + auto zero_scalar = cudf::numeric_scalar(0); + auto zero_literal = cudf::ast::literal(zero_scalar); + auto col_ref_0 = cudf::ast::column_name_reference("col0"); + auto mod_expression = cudf::ast::operation(cudf::ast::ast_operator::MOD, col_ref_0, two_literal); + auto filter_expression = + cudf::ast::operation(cudf::ast::ast_operator::EQUAL, mod_expression, zero_literal); + + auto const stream = cudf::get_default_stream(); + auto const mr = cudf::get_current_device_resource_ref(); + auto source_info = build_source_info(parquet_buffers); + + // Expected table read via the mainline multi-source parquet reader + auto const expected_options = cudf::io::parquet_reader_options::builder(source_info) + .filter(filter_expression) + .prepend_source_index_column(true) + .build(); + auto const expected = cudf::io::read_parquet(expected_options, stream, mr); + + // Hybrid scan multifile reader options + auto const options = cudf::io::parquet_reader_options::builder() + .filter(filter_expression) + .prepend_source_index_column(true) + .build(); + + auto inputs = multifile_inputs(source_info); + auto reader = + cudf::io::parquet::experimental::hybrid_scan_multifile{inputs.footer_byte_spans, options}; + + auto const row_group_indices = reader.all_row_groups(options); + auto row_mask = reader.build_all_true_row_mask(row_group_indices, stream, mr); + + auto filter_column_chunks = fetch_multisource_device_data( + inputs, reader.filter_column_chunks_byte_ranges(row_group_indices, options), stream, mr); + auto row_mask_view = row_mask->mutable_view(); + auto filter_result = reader.materialize_filter_columns(row_group_indices, + filter_column_chunks.flat_spans, + row_mask_view, + use_data_page_mask::NO, + options, + stream, + mr); + + // Filter result must contain the prepended source index column (col0 and src_idx) + ASSERT_EQ(filter_result.tbl->num_columns(), 2); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected.tbl->view(), filter_result.tbl->view()); +} + TEST_F(HybridScanMultifileTest, MaterializeStructs) { std::mt19937 gen(0xbaLL); diff --git a/cpp/tests/io/experimental/hybrid_scan_test.cpp b/cpp/tests/io/experimental/hybrid_scan_test.cpp index fe90fdebbb59..f7cd45b926c3 100644 --- a/cpp/tests/io/experimental/hybrid_scan_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_test.cpp @@ -878,148 +878,6 @@ TEST_F(HybridScanTest, StructChildFilterColumn) std::invalid_argument); } -TEST_F(HybridScanTest, PrependSourceIndexColumn) -{ - using T = uint32_t; - using cudf::io::parquet::experimental::use_data_page_mask; - - // Single source parquet file with columns col0 (filter), col1 and col2 (payload) - auto constexpr num_concat = 1; - auto [written_table, parquet_buffer] = create_parquet_with_stats(); - - // Filtering AST - col0 < 100 - auto literal_value = cudf::numeric_scalar(100); - auto literal = cudf::ast::literal(literal_value); - auto col_ref_0 = cudf::ast::column_name_reference("col0"); - auto filter_expression = cudf::ast::operation(cudf::ast::ast_operator::LESS, col_ref_0, literal); - - auto const stream = cudf::get_default_stream(); - auto const mr = cudf::get_current_device_resource_ref(); - - // Expected table - auto const expected = [&]() { - auto const options = - cudf::io::parquet_reader_options::builder( - cudf::io::source_info(cudf::host_span(parquet_buffer.data(), parquet_buffer.size()))) - .filter(filter_expression) - .prepend_source_index_column(true) - .build(); - return cudf::io::read_parquet(options, stream, mr).tbl; - }(); - - // Hybrid scan reader options - auto const options = cudf::io::parquet_reader_options::builder() - .filter(filter_expression) - .prepend_source_index_column(true) - .build(); - - auto datasource = cudf::io::datasource::create(cudf::host_span( - reinterpret_cast(parquet_buffer.data()), parquet_buffer.size())); - auto const footer_buffer = cudf::io::parquet::fetch_footer_to_host(*datasource); - - auto reader = - std::make_unique(*footer_buffer, options); - - auto const all_row_groups = reader->all_row_groups(options); - - auto const current_row_groups = cudf::host_span(all_row_groups); - auto const mid = current_row_groups.size() / 2; - auto const row_group_passes = - std::array{current_row_groups.subspan(0, mid), - current_row_groups.subspan(mid, current_row_groups.size() - mid)}; - auto const table_views = [](auto const& tables) { - auto views = std::vector{}; - std::transform(tables.begin(), tables.end(), std::back_inserter(views), [](auto const& t) { - return t->view(); - }); - return views; - }; - - // Two-step chunked read - { - auto filter_tables = std::vector>{}; - auto payload_tables = std::vector>{}; - - auto const read_pass = [&](auto const row_groups) { - auto row_mask = reader->build_all_true_row_mask(row_groups, stream, mr); - auto row_mask_view = row_mask->mutable_view(); - - auto const filter_byte_ranges = reader->filter_column_chunks_byte_ranges(row_groups, options); - auto [filter_bufs, filter_data, filter_tasks] = - cudf::io::parquet::fetch_byte_ranges_to_device_async( - *datasource, filter_byte_ranges, stream, mr); - filter_tasks.get(); - reader->setup_chunking_for_filter_columns(1024, - 10240, - row_groups, - row_mask_view, - use_data_page_mask::NO, - filter_data, - options, - stream, - mr); - while (reader->has_next_table_chunk()) { - filter_tables.push_back(reader->materialize_filter_columns_chunk(row_mask_view).tbl); - } - - auto const payload_byte_ranges = - reader->payload_column_chunks_byte_ranges(row_groups, options); - auto [payload_bufs, payload_data, payload_tasks] = - cudf::io::parquet::fetch_byte_ranges_to_device_async( - *datasource, payload_byte_ranges, stream, mr); - payload_tasks.get(); - - reader->setup_chunking_for_payload_columns(1024, - 10240, - row_groups, - row_mask_view, - use_data_page_mask::NO, - payload_data, - options, - stream, - mr); - while (reader->has_next_table_chunk()) { - payload_tables.push_back(reader->materialize_payload_columns_chunk(row_mask_view).tbl); - } - }; - - for (auto const& pass : row_group_passes) { - read_pass(pass); - } - - auto const filter_table = cudf::concatenate(table_views(filter_tables), stream, mr); - auto const payload_table = cudf::concatenate(table_views(payload_tables), stream, mr); - - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->select({0, 1}), filter_table->view()); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->select({2, 3}), payload_table->view()); - } - - // Single-step chunked read - { - auto all_tables = std::vector>{}; - - auto const read_pass = [&](cudf::host_span row_groups) { - auto const chunk_ranges = reader->all_column_chunks_byte_ranges(row_groups, options); - auto [bufs, col_data, tasks] = - cudf::io::parquet::fetch_byte_ranges_to_device_async(*datasource, chunk_ranges, stream, mr); - tasks.get(); - - reader->setup_chunking_for_all_columns( - 1024, 10240, row_groups, col_data, options, stream, mr); - while (reader->has_next_table_chunk()) { - all_tables.push_back(reader->materialize_all_columns_chunk().tbl); - } - }; - - for (auto const& pass : row_group_passes) { - read_pass(pass); - } - - auto const all_columns_table = cudf::concatenate(table_views(all_tables), stream, mr); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), all_columns_table->view()); - } -} - TEST_F(HybridScanTest, ChunkedReadRowMaskPerPass) { using T = uint32_t; diff --git a/cpp/tests/strings/contains_tests.cpp b/cpp/tests/strings/contains_tests.cpp index 89ed9756d00a..481d788ad3cb 100644 --- a/cpp/tests/strings/contains_tests.cpp +++ b/cpp/tests/strings/contains_tests.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/strings/findall_tests.cpp b/cpp/tests/strings/findall_tests.cpp index e4999a41bcc7..5076f6f7f50f 100644 --- a/cpp/tests/strings/findall_tests.cpp +++ b/cpp/tests/strings/findall_tests.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/strings/replace_regex_tests.cpp b/cpp/tests/strings/replace_regex_tests.cpp index 8422264d4c36..826d0e3cea41 100644 --- a/cpp/tests/strings/replace_regex_tests.cpp +++ b/cpp/tests/strings/replace_regex_tests.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ From 30e0ca0c92d92b1ca587911645cd2198c712273f Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Wed, 1 Jul 2026 20:58:47 +0000 Subject: [PATCH 06/15] Minor --- .../io/experimental/hybrid_scan_multifile_test.cpp | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp index 3fcecb1c2fd8..ca6ffecaa104 100644 --- a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp @@ -174,10 +174,11 @@ TEST_F(HybridScanMultifileTest, PrependSourceIndexColumn) cudf::io::table_input_metadata expected_metadata(table); expected_metadata.column_metadata[0].set_name("col0"); - std::vector> parquet_buffers(num_sources); - for (auto& parquet_buffer : parquet_buffers) { + // Write the table once and reference the same file for all sources + auto const parquet_filepath = temp_env->get_temp_filepath("PrependSourceIndexColumn.parquet"); + { auto out_opts = - cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&parquet_buffer}, table) + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{parquet_filepath}, table) .metadata(expected_metadata) .build(); cudf::io::write_parquet(out_opts); @@ -195,7 +196,9 @@ TEST_F(HybridScanMultifileTest, PrependSourceIndexColumn) auto const stream = cudf::get_default_stream(); auto const mr = cudf::get_current_device_resource_ref(); - auto source_info = build_source_info(parquet_buffers); + + auto const parquet_filepaths = std::vector(num_sources, parquet_filepath); + auto source_info = cudf::io::source_info(parquet_filepaths); // Expected table read via the mainline multi-source parquet reader auto const expected_options = cudf::io::parquet_reader_options::builder(source_info) From 52cc57494459b01d52e59e9d3db6aa78d48c6a35 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Wed, 1 Jul 2026 21:02:00 +0000 Subject: [PATCH 07/15] Minor --- cpp/src/io/parquet/reader_impl.hpp | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/cpp/src/io/parquet/reader_impl.hpp b/cpp/src/io/parquet/reader_impl.hpp index 32b96ec56b2a..ec7395376897 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -400,6 +400,18 @@ class reader_impl { return _subpass_page_mask ? *_subpass_page_mask : cudf::detail::hostdevice_span{}; } + /** + * @brief Offset the column references in `_expr_conv` by the number of columns prepended to + * the output + * + * @return Offsetted expression converter + */ + [[nodiscard]] inline offset_column_references compute_offset_filter() const + { + auto const num_prepended_cols = static_cast(_options.prepend_source_index_column); + return offset_column_references(_expr_conv.get_converted_expr(), num_prepended_cols); + } + /** * @brief Calculate the number of rows read from each source in the output chunk * @@ -419,18 +431,6 @@ class reader_impl { void prepend_source_index_column(std::span num_rows_per_source, std::vector>& out_columns); - /** - * @brief Offset the column references in `_expr_conv` by the number of columns prepended to - * the output - * - * @return Offsetted expression converter - */ - [[nodiscard]] offset_column_references compute_offset_filter() const - { - auto const num_prepended_cols = static_cast(_options.prepend_source_index_column); - return offset_column_references(_expr_conv.get_converted_expr(), num_prepended_cols); - } - /** * @brief Computes the names of columns to be read from the file, if specified. * From 496fa6f38be524a68096a2586cafa6aea7814267 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Thu, 9 Jul 2026 19:24:18 +0000 Subject: [PATCH 08/15] Update --- .../parquet/experimental/hybrid_scan_impl.cpp | 27 ++++++++++++++++--- cpp/src/io/parquet/reader_impl.cpp | 5 +--- cpp/src/io/parquet/reader_impl.hpp | 3 ++- cpp/tests/strings/contains_tests.cpp | 2 +- cpp/tests/strings/findall_tests.cpp | 2 +- cpp/tests/strings/replace_regex_tests.cpp | 2 +- 6 files changed, 29 insertions(+), 12 deletions(-) diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index 032554a2940a..56a53721f82d 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -902,6 +902,7 @@ void hybrid_scan_reader_impl::reset_internal_state() _options.use_jit_filter = false; _options.case_sensitive_names = true; _options.prepend_source_index_column = false; + _options.prepend_row_index_column = false; _num_sources = 0; _input_pass_read_limit = 0; @@ -923,6 +924,7 @@ void hybrid_scan_reader_impl::initialize_column_selection_options( _options.use_jit_filter = options.is_enabled_use_jit_filter(); _options.case_sensitive_names = options.is_enabled_case_sensitive_names(); _options.prepend_source_index_column = options.is_enabled_prepend_source_index_column(); + _options.prepend_row_index_column = options.is_enabled_prepend_row_index_column(); _use_pandas_metadata = options.is_enabled_use_pandas_metadata(); } @@ -1112,6 +1114,13 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output( _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] + : cudf::io::parquet::detail::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; @@ -1124,10 +1133,20 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output( apply_decimal_width_cast(out_columns); - // Prepend the source index column to filter columns only - if (_options.prepend_source_index_column and - read_columns_mode == read_columns_mode::FILTER_COLUMNS) { - prepend_source_index_column(out_metadata.num_rows_per_source, out_columns); + // 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}); + } } // Compute the final filter expression incorporating any column reference offsets in _expr_conv diff --git a/cpp/src/io/parquet/reader_impl.cpp b/cpp/src/io/parquet/reader_impl.cpp index 1155628159bf..b4e80c13656e 100644 --- a/cpp/src/io/parquet/reader_impl.cpp +++ b/cpp/src/io/parquet/reader_impl.cpp @@ -911,10 +911,7 @@ table_with_metadata reader_impl::finalize_output(read_mode mode, } // Offset column references in `_expr_conv` by the number of prepended columns - 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 = compute_offset_filter(); auto const final_filter_expr = final_filter.get_converted_expr(); // check if the output filter AST expression (= _expr_conv.get_converted_expr()) exists diff --git a/cpp/src/io/parquet/reader_impl.hpp b/cpp/src/io/parquet/reader_impl.hpp index 87205510c686..1b415e89c21d 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -408,7 +408,8 @@ class reader_impl { */ [[nodiscard]] inline offset_column_references compute_offset_filter() const { - 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); return offset_column_references(_expr_conv.get_converted_expr(), num_prepended_cols); } diff --git a/cpp/tests/strings/contains_tests.cpp b/cpp/tests/strings/contains_tests.cpp index 481d788ad3cb..89ed9756d00a 100644 --- a/cpp/tests/strings/contains_tests.cpp +++ b/cpp/tests/strings/contains_tests.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/strings/findall_tests.cpp b/cpp/tests/strings/findall_tests.cpp index 5076f6f7f50f..e4999a41bcc7 100644 --- a/cpp/tests/strings/findall_tests.cpp +++ b/cpp/tests/strings/findall_tests.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/tests/strings/replace_regex_tests.cpp b/cpp/tests/strings/replace_regex_tests.cpp index 826d0e3cea41..8422264d4c36 100644 --- a/cpp/tests/strings/replace_regex_tests.cpp +++ b/cpp/tests/strings/replace_regex_tests.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ From ccadf62249fb942dbdc9b7315b8a839334f256c5 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Thu, 9 Jul 2026 19:37:15 +0000 Subject: [PATCH 09/15] Minor bug fix --- cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index 56a53721f82d..06e4992933e7 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -1133,8 +1133,8 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output( apply_decimal_width_cast(out_columns); - // Prepend the source and row index columns if requested - { + // Prepend the source and row index columns to filter columns only + if (read_columns_mode == read_columns_mode::FILTER_COLUMNS) { 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(), @@ -1149,10 +1149,6 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output( } } - // Compute the final filter expression incorporating any column reference offsets in _expr_conv - auto const final_filter = compute_offset_filter(); - auto const final_filter_expr = final_filter.get_converted_expr(); - // Create a table from the output columns. auto read_table = std::make_unique
(std::move(out_columns)); @@ -1177,6 +1173,10 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output( if constexpr (std::is_same_v) { CUDF_EXPECTS(read_columns_mode == read_columns_mode::FILTER_COLUMNS, "Invalid read mode"); + // Compute the final filter expression incorporating any column reference offsets in _expr_conv + auto const final_filter = compute_offset_filter(); + auto const final_filter_expr = final_filter.get_converted_expr(); + auto final_row_mask = cudf::detail::compute_column(*read_table, final_filter_expr.value().get(), _stream, From db8dfe356e1abb70cd2f6eb9270b721e079fb23a Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Thu, 9 Jul 2026 21:07:25 +0000 Subject: [PATCH 10/15] Supply stream and mr to column synthesizers --- cpp/src/io/parquet/reader_impl.cpp | 8 ++-- cpp/src/io/parquet/reader_impl.hpp | 11 +++++- cpp/src/io/parquet/reader_impl_preprocess.cu | 39 +++++++++++--------- 3 files changed, 36 insertions(+), 22 deletions(-) diff --git a/cpp/src/io/parquet/reader_impl.cpp b/cpp/src/io/parquet/reader_impl.cpp index b4e80c13656e..e46e3f22b3d6 100644 --- a/cpp/src/io/parquet/reader_impl.cpp +++ b/cpp/src/io/parquet/reader_impl.cpp @@ -897,13 +897,15 @@ 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) { - out_columns.emplace(out_columns.begin(), synthesize_row_index_column(read_info)); + out_columns.emplace(out_columns.begin(), + synthesize_row_index_column(read_info, _stream, _mr)); 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_columns.emplace( + out_columns.begin(), + synthesize_source_index_column(out_metadata.num_rows_per_source, _stream, _mr)); 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.hpp b/cpp/src/io/parquet/reader_impl.hpp index 1b415e89c21d..ecf95a90eb25 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -427,10 +427,14 @@ class reader_impl { * @brief Synthesize source index column * * @param num_rows_per_source Number of rows per parquet source + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource to use for device memory allocation * @return Synthesized source index column */ [[nodiscard]] std::unique_ptr synthesize_source_index_column( - std::span num_rows_per_source); + std::span num_rows_per_source, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); /** * @brief Synthesize file-local row index column @@ -440,9 +444,12 @@ class reader_impl { * * @param read_info Row range of the output chunk relative to the first row of the first * selected row group + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource to use for device memory allocation * @return Synthesized row index column */ - [[nodiscard]] std::unique_ptr synthesize_row_index_column(row_range const& read_info); + [[nodiscard]] std::unique_ptr synthesize_row_index_column( + row_range const& read_info, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); /** * @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 bc0540e5a919..7634bb90cdcc 100644 --- a/cpp/src/io/parquet/reader_impl_preprocess.cu +++ b/cpp/src/io/parquet/reader_impl_preprocess.cu @@ -1152,7 +1152,9 @@ struct map_global_to_local_row_index { } // namespace -std::unique_ptr reader_impl::synthesize_row_index_column(row_range const& read_info) +std::unique_ptr reader_impl::synthesize_row_index_column(row_range const& read_info, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { using column_type = size_t; @@ -1161,16 +1163,16 @@ std::unique_ptr reader_impl::synthesize_row_index_column(row_range const } // Allocate column data vector - auto col_data = rmm::device_uvector(read_info.num_rows, _stream, _mr); + 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); + 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); + 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); @@ -1178,9 +1180,9 @@ std::unique_ptr reader_impl::synthesize_row_index_column(row_range const // 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()); + 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()); + 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( @@ -1189,15 +1191,17 @@ std::unique_ptr reader_impl::synthesize_row_index_column(row_range const 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(); + stream.value())); + stream.synchronize(); } - return std::make_unique(std::move(col_data), rmm::device_buffer{}, 0); + return std::make_unique(std::move(col_data), rmm::device_buffer{0, stream, mr}, 0); } std::unique_ptr reader_impl::synthesize_source_index_column( - std::span num_rows_per_source) + std::span num_rows_per_source, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { using column_type = cudf::size_type; @@ -1211,12 +1215,13 @@ std::unique_ptr reader_impl::synthesize_source_index_column( // Single source if (num_sources == 1) { - auto const scalar = cudf::numeric_scalar(0, true, _stream, _mr); - return cudf::make_column_from_scalar(scalar, num_rows, _stream, _mr); + auto const scalar = + cudf::numeric_scalar(0, true, stream, cudf::get_current_device_resource_ref()); + return cudf::make_column_from_scalar(scalar, num_rows, stream, mr); } // Allocate column data vector - auto col_data = rmm::device_uvector(num_rows, _stream, _mr); + auto col_data = rmm::device_uvector(num_rows, stream, mr); // Label each output row with its source index via segment boundaries. { @@ -1228,13 +1233,13 @@ std::unique_ptr reader_impl::synthesize_source_index_column( 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, cudf::get_current_device_resource_ref()); + 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(); + 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); + return std::make_unique(std::move(col_data), rmm::device_buffer{0, stream, mr}, 0); } } // namespace cudf::io::parquet::detail From dae280c25784a5f9ba631f9a9849a809f3cb2c9a Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Fri, 10 Jul 2026 18:44:51 +0000 Subject: [PATCH 11/15] Sync with updated column synthesizers --- cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index 06e4992933e7..1ddc3e9ec847 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -1136,13 +1136,15 @@ table_with_metadata hybrid_scan_reader_impl::finalize_output( // Prepend the source and row index columns to filter columns only if (read_columns_mode == read_columns_mode::FILTER_COLUMNS) { if (_options.prepend_row_index_column) { - out_columns.emplace(out_columns.begin(), synthesize_row_index_column(read_info)); + out_columns.emplace(out_columns.begin(), + synthesize_row_index_column(read_info, _stream, _mr)); 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_columns.emplace( + out_columns.begin(), + synthesize_source_index_column(out_metadata.num_rows_per_source, _stream, _mr)); out_metadata.schema_info.emplace( out_metadata.schema_info.begin(), column_name_info{.name = "source_index", .is_nullable = false}); From 7ce930b0f81859e3b8ca3f88d110db24eb4cc536 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:28:51 +0000 Subject: [PATCH 12/15] Minor --- cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp | 2 -- 1 file changed, 2 deletions(-) diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index f0c51a7c7ce4..7fda93d76c31 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -898,8 +898,6 @@ void hybrid_scan_reader_impl::initialize_reader_config(parquet_reader_options co _options.case_sensitive_names = options.is_enabled_case_sensitive_names(); _options.prepend_source_index_column = options.is_enabled_prepend_source_index_column(); _options.prepend_row_index_column = options.is_enabled_prepend_row_index_column(); - - _use_pandas_metadata = options.is_enabled_use_pandas_metadata(); } void hybrid_scan_reader_impl::initialize_options(parquet_reader_options const& options, From 5d718590af6bf7e24b00961c38684a0984dcc471 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:36:11 +0000 Subject: [PATCH 13/15] Improve test coverage --- .../hybrid_scan_multifile_test.cpp | 31 +++++++++++++++---- 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp index ca6ffecaa104..fd7db3d35b0f 100644 --- a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp @@ -159,7 +159,7 @@ TEST_F(HybridScanMultifileTest, MaterializeListsOfStrings) test_hybrid_scan_multifile({col0, *col1, *col2, *col3, *col4}, false); } -TEST_F(HybridScanMultifileTest, PrependSourceIndexColumn) +TEST_F(HybridScanMultifileTest, PrependIndexColumns) { using T = int32_t; using cudf::io::parquet::experimental::use_data_page_mask; @@ -175,7 +175,7 @@ TEST_F(HybridScanMultifileTest, PrependSourceIndexColumn) expected_metadata.column_metadata[0].set_name("col0"); // Write the table once and reference the same file for all sources - auto const parquet_filepath = temp_env->get_temp_filepath("PrependSourceIndexColumn.parquet"); + auto const parquet_filepath = temp_env->get_temp_filepath("PrependIndexColumns.parquet"); { auto out_opts = cudf::io::parquet_writer_options::builder(cudf::io::sink_info{parquet_filepath}, table) @@ -204,13 +204,14 @@ TEST_F(HybridScanMultifileTest, PrependSourceIndexColumn) auto const expected_options = cudf::io::parquet_reader_options::builder(source_info) .filter(filter_expression) .prepend_source_index_column(true) + .prepend_row_index_column(true) .build(); - auto const expected = cudf::io::read_parquet(expected_options, stream, mr); // Hybrid scan multifile reader options auto const options = cudf::io::parquet_reader_options::builder() .filter(filter_expression) .prepend_source_index_column(true) + .prepend_row_index_column(true) .build(); auto inputs = multifile_inputs(source_info); @@ -231,9 +232,27 @@ TEST_F(HybridScanMultifileTest, PrependSourceIndexColumn) stream, mr); - // Filter result must contain the prepended source index column (col0 and src_idx) - ASSERT_EQ(filter_result.tbl->num_columns(), 2); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected.tbl->view(), filter_result.tbl->view()); + // Build expected table with source and row index columns + auto const source_index = + cudf::detail::make_counting_transform_iterator(0, [](cudf::size_type i) { return i / 5; }); + auto const expected_source_index = cudf::test::fixed_width_column_wrapper( + source_index, source_index + num_sources * 5); + + auto const row_index = cudf::detail::make_counting_transform_iterator( + 0, [](cudf::size_type i) -> size_t { return (i % 5) * 2; }); + auto const expected_row_index = + cudf::test::fixed_width_column_wrapper(row_index, row_index + num_sources * 5); + + auto const filtered_values = cudf::detail::make_counting_transform_iterator( + 0, [](cudf::size_type i) -> T { return static_cast((i % 5) * 2); }); + auto const expected_values = + cudf::test::fixed_width_column_wrapper(filtered_values, filtered_values + num_sources * 5); + + auto const expected_table = + cudf::table_view{{expected_source_index, expected_row_index, expected_values}}; + + ASSERT_EQ(filter_result.tbl->num_columns(), 3); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table, filter_result.tbl->view()); } TEST_F(HybridScanMultifileTest, MaterializeStructs) From 46b31bfb113c9b6a1eee97f2f5ae26f4e4061f8d Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:50:40 +0000 Subject: [PATCH 14/15] Test coverage --- .../hybrid_scan_multifile_test.cpp | 91 ++++++++++--------- 1 file changed, 50 insertions(+), 41 deletions(-) diff --git a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp index fd7db3d35b0f..f297d000b77b 100644 --- a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp @@ -169,10 +169,11 @@ TEST_F(HybridScanMultifileTest, PrependIndexColumns) auto constexpr num_sources = 3; auto values = cuda::counting_iterator{0}; cudf::test::fixed_width_column_wrapper col0(values, values + num_rows); - auto const table = cudf::table_view{{col0}}; + auto const table = cudf::table_view{{col0, col0}}; cudf::io::table_input_metadata expected_metadata(table); expected_metadata.column_metadata[0].set_name("col0"); + expected_metadata.column_metadata[1].set_name("col1"); // Write the table once and reference the same file for all sources auto const parquet_filepath = temp_env->get_temp_filepath("PrependIndexColumns.parquet"); @@ -194,44 +195,6 @@ TEST_F(HybridScanMultifileTest, PrependIndexColumns) auto filter_expression = cudf::ast::operation(cudf::ast::ast_operator::EQUAL, mod_expression, zero_literal); - auto const stream = cudf::get_default_stream(); - auto const mr = cudf::get_current_device_resource_ref(); - - auto const parquet_filepaths = std::vector(num_sources, parquet_filepath); - auto source_info = cudf::io::source_info(parquet_filepaths); - - // Expected table read via the mainline multi-source parquet reader - auto const expected_options = cudf::io::parquet_reader_options::builder(source_info) - .filter(filter_expression) - .prepend_source_index_column(true) - .prepend_row_index_column(true) - .build(); - - // Hybrid scan multifile reader options - auto const options = cudf::io::parquet_reader_options::builder() - .filter(filter_expression) - .prepend_source_index_column(true) - .prepend_row_index_column(true) - .build(); - - auto inputs = multifile_inputs(source_info); - auto reader = - cudf::io::parquet::experimental::hybrid_scan_multifile{inputs.footer_byte_spans, options}; - - auto const row_group_indices = reader.all_row_groups(options); - auto row_mask = reader.build_all_true_row_mask(row_group_indices, stream, mr); - - auto filter_column_chunks = fetch_multisource_device_data( - inputs, reader.filter_column_chunks_byte_ranges(row_group_indices, options), stream, mr); - auto row_mask_view = row_mask->mutable_view(); - auto filter_result = reader.materialize_filter_columns(row_group_indices, - filter_column_chunks.flat_spans, - row_mask_view, - use_data_page_mask::NO, - options, - stream, - mr); - // Build expected table with source and row index columns auto const source_index = cudf::detail::make_counting_transform_iterator(0, [](cudf::size_type i) { return i / 5; }); @@ -251,8 +214,54 @@ TEST_F(HybridScanMultifileTest, PrependIndexColumns) auto const expected_table = cudf::table_view{{expected_source_index, expected_row_index, expected_values}}; - ASSERT_EQ(filter_result.tbl->num_columns(), 3); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table, filter_result.tbl->view()); + // Hybrid scan multifile reader options + auto const options = cudf::io::parquet_reader_options::builder() + .filter(filter_expression) + .prepend_source_index_column(true) + .prepend_row_index_column(true) + .build(); + + auto const parquet_filepaths = std::vector(num_sources, parquet_filepath); + auto inputs = multifile_inputs(cudf::io::source_info(parquet_filepaths)); + auto reader = + cudf::io::parquet::experimental::hybrid_scan_multifile{inputs.footer_byte_spans, options}; + + auto const stream = cudf::get_default_stream(); + auto const mr = cudf::get_current_device_resource_ref(); + + auto const row_group_indices = reader.all_row_groups(options); + auto row_mask = reader.build_all_true_row_mask(row_group_indices, stream, mr); + + // Materialize filter column prepended with index columns + { + auto filter_column_chunks = fetch_multisource_device_data( + inputs, reader.filter_column_chunks_byte_ranges(row_group_indices, options), stream, mr); + auto row_mask_view = row_mask->mutable_view(); + auto filter_result = reader.materialize_filter_columns(row_group_indices, + filter_column_chunks.flat_spans, + row_mask_view, + use_data_page_mask::NO, + options, + stream, + mr); + + ASSERT_EQ(filter_result.tbl->num_columns(), 3); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table, filter_result.tbl->view()); + + // Materialize payload column (no prepended index columns) + auto payload_column_chunks = fetch_multisource_device_data( + inputs, reader.payload_column_chunks_byte_ranges(row_group_indices, options), stream, mr); + auto payload_result = reader.materialize_payload_columns(row_group_indices, + payload_column_chunks.flat_spans, + row_mask->view(), + use_data_page_mask::NO, + options, + stream, + mr); + ASSERT_EQ(payload_result.tbl->num_columns(), 1); + // col1 (payload) must be identical to col0 (filter) with the same row_mask + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table.select({2}), payload_result.tbl->view()); + } } TEST_F(HybridScanMultifileTest, MaterializeStructs) From ded67a222a4d4a5510c6fc1c33a22018f907c8cc Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Mon, 13 Jul 2026 19:51:25 +0000 Subject: [PATCH 15/15] whitespace changes --- .../hybrid_scan_multifile_test.cpp | 48 +++++++++---------- 1 file changed, 23 insertions(+), 25 deletions(-) diff --git a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp index f297d000b77b..9a4d9daa2c97 100644 --- a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp @@ -233,35 +233,33 @@ TEST_F(HybridScanMultifileTest, PrependIndexColumns) auto row_mask = reader.build_all_true_row_mask(row_group_indices, stream, mr); // Materialize filter column prepended with index columns - { - auto filter_column_chunks = fetch_multisource_device_data( - inputs, reader.filter_column_chunks_byte_ranges(row_group_indices, options), stream, mr); - auto row_mask_view = row_mask->mutable_view(); - auto filter_result = reader.materialize_filter_columns(row_group_indices, - filter_column_chunks.flat_spans, - row_mask_view, + auto filter_column_chunks = fetch_multisource_device_data( + inputs, reader.filter_column_chunks_byte_ranges(row_group_indices, options), stream, mr); + auto row_mask_view = row_mask->mutable_view(); + auto filter_result = reader.materialize_filter_columns(row_group_indices, + filter_column_chunks.flat_spans, + row_mask_view, + use_data_page_mask::NO, + options, + stream, + mr); + + ASSERT_EQ(filter_result.tbl->num_columns(), 3); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table, filter_result.tbl->view()); + + // Materialize payload column (no prepended index columns) + auto payload_column_chunks = fetch_multisource_device_data( + inputs, reader.payload_column_chunks_byte_ranges(row_group_indices, options), stream, mr); + auto payload_result = reader.materialize_payload_columns(row_group_indices, + payload_column_chunks.flat_spans, + row_mask->view(), use_data_page_mask::NO, options, stream, mr); - - ASSERT_EQ(filter_result.tbl->num_columns(), 3); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table, filter_result.tbl->view()); - - // Materialize payload column (no prepended index columns) - auto payload_column_chunks = fetch_multisource_device_data( - inputs, reader.payload_column_chunks_byte_ranges(row_group_indices, options), stream, mr); - auto payload_result = reader.materialize_payload_columns(row_group_indices, - payload_column_chunks.flat_spans, - row_mask->view(), - use_data_page_mask::NO, - options, - stream, - mr); - ASSERT_EQ(payload_result.tbl->num_columns(), 1); - // col1 (payload) must be identical to col0 (filter) with the same row_mask - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table.select({2}), payload_result.tbl->view()); - } + ASSERT_EQ(payload_result.tbl->num_columns(), 1); + // col1 (payload) must be identical to col0 (filter) with the same row_mask + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_table.select({2}), payload_result.tbl->view()); } TEST_F(HybridScanMultifileTest, MaterializeStructs)