diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 63c292646773..d3576d2a6b5a 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -627,6 +627,7 @@ add_library( src/io/parquet/experimental/hybrid_scan_chunking.cu src/io/parquet/experimental/hybrid_scan_helpers.cpp src/io/parquet/experimental/hybrid_scan_impl.cpp + src/io/parquet/experimental/hybrid_scan_multifile.cpp src/io/parquet/experimental/hybrid_scan_preprocess.cu src/io/parquet/experimental/page_index_filter.cu src/io/parquet/experimental/page_index_filter_utils.cu diff --git a/cpp/include/cudf/io/experimental/hybrid_scan.hpp b/cpp/include/cudf/io/experimental/hybrid_scan.hpp index 0f7ead8a3dfd..980ab9644d3b 100644 --- a/cpp/include/cudf/io/experimental/hybrid_scan.hpp +++ b/cpp/include/cudf/io/experimental/hybrid_scan.hpp @@ -19,15 +19,13 @@ #include #include -namespace CUDF_EXPORT cudf { -namespace io::parquet::experimental::detail { +namespace cudf::io::parquet::experimental::detail { /** * @brief Internal experimental Parquet reader optimized for highly selective filters, called a * Hybrid Scan operation. */ class hybrid_scan_reader_impl; -} // namespace io::parquet::experimental::detail -} // namespace CUDF_EXPORT cudf +} // namespace cudf::io::parquet::experimental::detail //! Using `byte_range_info` from cudf::io::text using cudf::io::text::byte_range_info; @@ -344,7 +342,7 @@ class hybrid_scan_reader { * @param row_group_indices Input row groups indices * @return Total number of top-level rows in the row groups */ - [[nodiscard]] size_type total_rows_in_row_groups( + [[nodiscard]] std::size_t total_rows_in_row_groups( cudf::host_span row_group_indices) const; /** diff --git a/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp b/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp new file mode 100644 index 000000000000..70ff792660bf --- /dev/null +++ b/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp @@ -0,0 +1,133 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include + +namespace cudf::io::parquet::experimental::detail { +/** + * @brief Internal experimental Parquet reader optimized for highly selective filters, called a + * Hybrid Scan operation. + */ +class hybrid_scan_reader_impl; +} // namespace cudf::io::parquet::experimental::detail + +//! Using `byte_range_info` from cudf::io::text +using cudf::io::text::byte_range_info; + +namespace CUDF_EXPORT cudf { +namespace io::parquet::experimental { +/** + * @addtogroup io_readers + * @{ + * @file + */ + +/** + * @brief Multi-file variant of the experimental Hybrid Scan Parquet reader + * + * Vectorizes `hybrid_scan_reader` APIs to support multiple Parquet sources. Inputs and outputs are + * indexed by source order except for the row mask which is a single BOOL8 column spanning all rows + * from all sources concatenated in source order, then row-group order within a source. + * + * @note Detailed usage documentation will be added once all APIs are in place. This reader will + * eventually move to `hybrid_scan.hpp` and the existing single-file reader (`hybrid_scan_reader`) + * will become its subclass. Only keeping this separate here for now to reduce noise. + */ +class hybrid_scan_multifile { + public: + /** + * @brief Constructor for the multi-file experimental Parquet reader + * + * @param footer_bytes Host span of Parquet file footer byte spans, one per source + * @param options Parquet reader options + */ + explicit hybrid_scan_multifile(cudf::host_span const> footer_bytes, + parquet_reader_options const& options); + + /** + * @brief Constructor for the multi-file experimental Parquet reader + * + * @param parquet_metadata Host span of pre-populated Parquet file metadata, one per source + * @param options Parquet reader options + */ + explicit hybrid_scan_multifile(cudf::host_span parquet_metadata, + parquet_reader_options const& options); + + /** + * @brief Destructor for the multi-file experimental Parquet reader + */ + ~hybrid_scan_multifile(); + + /** + * @brief Get parquet metadatas for all sources + * + * @return Vector of parquet metadata, one per source + */ + [[nodiscard]] std::vector parquet_metadatas() const; + + /** + * @brief Get byte ranges of the page index for all sources + * + * @return Vector of page index byte ranges, one per source + */ + [[nodiscard]] std::vector page_index_byte_ranges() const; + + /** + * @brief Setup the per-source page index within each Parquet file metadata + * + * @param page_index_bytes Host span of Parquet page index buffer bytes, one per source + */ + void setup_page_indexes( + cudf::host_span const> page_index_bytes) const; + + /** + * @brief Get all available per-source row group indices from the parquet files + * + * @param options Parquet reader options + * @return Vector of row group indices, one inner vector per source + */ + [[nodiscard]] std::vector> all_row_groups( + parquet_reader_options const& options) const; + + /** + * @brief Get the total number of top-level rows in the per-source row groups + * + * @param row_group_indices Input per-source row group indices (one inner vector per source) + * @return Total number of top-level rows across all sources + */ + [[nodiscard]] size_type total_rows_in_row_groups( + cudf::host_span const> row_group_indices) const; + + /** + * @brief Resets the current column selection + * + * Resets the current column selection state forcing column re-selection in subsequent filter, + * byte range, setup chunking and materialization APIs. This is useful if the filter expression + * has been cascaded (and-ed) to include new columns. + */ + void reset_column_selection() const; + + private: + std::unique_ptr _impl; +}; + +/** @} */ // end of group + +} // namespace io::parquet::experimental +} // namespace CUDF_EXPORT cudf diff --git a/cpp/src/io/parquet/experimental/hybrid_scan.cpp b/cpp/src/io/parquet/experimental/hybrid_scan.cpp index b6243fafe60a..868813a1b4ed 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan.cpp @@ -15,13 +15,15 @@ namespace cudf::io::parquet::experimental { hybrid_scan_reader::hybrid_scan_reader(cudf::host_span footer_bytes, parquet_reader_options const& options) - : _impl{std::make_unique(footer_bytes, options)} + : _impl{std::make_unique( + std::vector>{footer_bytes}, options)} { } hybrid_scan_reader::hybrid_scan_reader(FileMetaData const& parquet_metadata, parquet_reader_options const& options) - : _impl{std::make_unique(parquet_metadata, options)} + : _impl{std::make_unique( + std::vector{parquet_metadata}, options)} { } @@ -29,19 +31,18 @@ hybrid_scan_reader::~hybrid_scan_reader() = default; [[nodiscard]] text::byte_range_info hybrid_scan_reader::page_index_byte_range() const { - return _impl->page_index_byte_range(); + return _impl->page_index_byte_ranges().front(); } [[nodiscard]] FileMetaData hybrid_scan_reader::parquet_metadata() const { - return _impl->parquet_metadata(); + return _impl->parquet_metadatas().front(); } void hybrid_scan_reader::setup_page_index(cudf::host_span page_index_bytes) const { CUDF_FUNC_RANGE(); - - return _impl->setup_page_index(page_index_bytes); + return _impl->setup_page_indexes(std::vector>{page_index_bytes}); } std::vector hybrid_scan_reader::all_row_groups( @@ -53,10 +54,10 @@ std::vector hybrid_scan_reader::all_row_groups( // If row groups are specified in parquet reader options, return them as is if (options.get_row_groups().size() == 1) { return options.get_row_groups().front(); } - return _impl->all_row_groups(options); + return _impl->all_row_groups(options).front(); } -size_type hybrid_scan_reader::total_rows_in_row_groups( +std::size_t hybrid_scan_reader::total_rows_in_row_groups( cudf::host_span row_group_indices) const { if (row_group_indices.empty()) { return 0; } diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_helpers.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_helpers.cpp index f6e6985ea4ea..626ac249b1bd 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_helpers.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_helpers.cpp @@ -72,23 +72,34 @@ metadata::metadata(cudf::host_span footer_bytes) sanitize_schema(); } -aggregate_reader_metadata::aggregate_reader_metadata(FileMetaData const& parquet_metadata, - bool use_arrow_schema, - bool has_cols_from_mismatched_srcs) +aggregate_reader_metadata::aggregate_reader_metadata( + cudf::host_span const> footer_bytes, + bool use_arrow_schema, + bool has_cols_from_mismatched_srcs) : aggregate_reader_metadata_base(host_span const>{}, false, false) { - // Just copy over the FileMetaData struct to the internal metadata struct - per_file_metadata.emplace_back(metadata{parquet_metadata}); + CUDF_EXPECTS(not footer_bytes.empty(), "At least one source must be provided"); + per_file_metadata.reserve(footer_bytes.size()); + std::transform(footer_bytes.begin(), + footer_bytes.end(), + std::back_inserter(per_file_metadata), + [](auto const& fb) { return metadata{fb}; }); initialize_internals(use_arrow_schema, has_cols_from_mismatched_srcs); } -aggregate_reader_metadata::aggregate_reader_metadata(cudf::host_span footer_bytes, - bool use_arrow_schema, - bool has_cols_from_mismatched_srcs) +aggregate_reader_metadata::aggregate_reader_metadata( + cudf::host_span parquet_metadatas, + bool use_arrow_schema, + bool has_cols_from_mismatched_srcs) : aggregate_reader_metadata_base(host_span const>{}, false, false) { - // Re-initialize internal variables here as base class was initialized without a source - per_file_metadata.emplace_back(metadata{footer_bytes}); + CUDF_EXPECTS(not parquet_metadatas.empty(), "At least one source must be provided"); + per_file_metadata.reserve(parquet_metadatas.size()); + // Just copy over the FileMetaData structs to the internal metadata structs + std::transform(parquet_metadatas.begin(), + parquet_metadatas.end(), + std::back_inserter(per_file_metadata), + [](auto const& parquet_metadata) { return metadata{parquet_metadata}; }); initialize_internals(use_arrow_schema, has_cols_from_mismatched_srcs); } @@ -102,13 +113,15 @@ void aggregate_reader_metadata::initialize_internals(bool use_arrow_schema, // Force all non-nullable (REQUIRED) columns to be nullable without modifying REPEATED columns to // preserve list structures - auto& schema = per_file_metadata.front().schema; - std::for_each(schema.begin() + 1, schema.end(), [](auto& col) { - // TODO: Store information of whichever column schema we modified here and restore it to - // `REQUIRED` if we end up not pruning any pages out of it - if (col.repetition_type == FieldRepetitionType::REQUIRED) { - col.repetition_type = FieldRepetitionType::OPTIONAL; - } + std::for_each(per_file_metadata.begin(), per_file_metadata.end(), [](auto& pfm) { + auto& schema = pfm.schema; + std::for_each(schema.begin() + 1, schema.end(), [](auto& col) { + // TODO: Store information of whichever column schema we modified here and restore it to + // `REQUIRED` if we end up not pruning any pages out of it + if (col.repetition_type == FieldRepetitionType::REQUIRED) { + col.repetition_type = FieldRepetitionType::OPTIONAL; + } + }); }); // Collect and apply arrow:schema from Parquet's key value metadata section @@ -122,74 +135,117 @@ void aggregate_reader_metadata::initialize_internals(bool use_arrow_schema, } } -text::byte_range_info aggregate_reader_metadata::page_index_byte_range() const +std::vector aggregate_reader_metadata::page_index_byte_ranges() const { - auto& schema = per_file_metadata.front(); - auto& row_groups = schema.row_groups; - - if (row_groups.size() and row_groups.front().columns.size()) { - auto const min_offset = schema.row_groups.front().columns.front().column_index_offset; - auto const& last_col = schema.row_groups.back().columns.back(); - auto const max_offset = last_col.offset_index_offset + last_col.offset_index_length; - return {min_offset, (max_offset - min_offset)}; - } + std::vector page_index_byte_ranges; + std::transform(per_file_metadata.begin(), + per_file_metadata.end(), + std::back_inserter(page_index_byte_ranges), + [](auto const& file_metadata) -> text::byte_range_info { + auto const& row_groups = file_metadata.row_groups; + if (row_groups.empty() or row_groups.front().columns.empty()) { return {}; } + + auto const min_offset = row_groups.front().columns.front().column_index_offset; + auto const& last_col = row_groups.back().columns.back(); + auto const max_offset = + last_col.offset_index_offset + last_col.offset_index_length; + + if (max_offset <= min_offset) { return {}; } + return {min_offset, max_offset - min_offset}; + }); - return {}; + return page_index_byte_ranges; } -FileMetaData aggregate_reader_metadata::parquet_metadata() const +std::vector aggregate_reader_metadata::parquet_metadatas() const { - return per_file_metadata.front(); + return {per_file_metadata.begin(), per_file_metadata.end()}; } -void aggregate_reader_metadata::setup_page_index(cudf::host_span page_index_bytes) +void aggregate_reader_metadata::setup_page_indexes( + cudf::host_span const> page_index_bytes) { - // Return early if empty page index buffer span - if (page_index_bytes.empty()) { - CUDF_LOG_WARN("Hybrid scan reader encountered empty page index buffer"); - return; - } - - // Get the file metadata and setup the page index - auto& file_metadata = per_file_metadata.front(); - auto const& row_groups = file_metadata.row_groups; - - // Check for empty parquet file - CUDF_EXPECTS(not row_groups.empty() and not row_groups.front().columns.empty(), - "No column chunks in Parquet schema to read page index for"); + CUDF_EXPECTS(page_index_bytes.size() == per_file_metadata.size(), + "Page index byte span count must equal the number of sources"); + + auto iter = cuda::zip_iterator(page_index_bytes.begin(), per_file_metadata.begin()); + std::for_each(iter, iter + page_index_bytes.size(), [&](auto const& pair) { + // Get the page index bytes and file metadata + auto const& [pgidx_bytes, file_metadata] = pair; + auto const& row_groups = file_metadata.row_groups; + + // Return early if empty page index buffer span + if (pgidx_bytes.empty()) { return; } + + // Check for empty parquet file + CUDF_EXPECTS(not row_groups.empty() and not row_groups.front().columns.empty(), + "No column chunks in Parquet schema to read page index for"); + + // Set the first ColumnChunk's offset of ColumnIndex as the adjusted zero offset + int64_t const min_offset = row_groups.front().columns.front().column_index_offset; + + // Check if the page index buffer is valid + { + auto const& last_col = row_groups.back().columns.back(); + auto const max_offset = last_col.offset_index_offset + last_col.offset_index_length; + CUDF_EXPECTS(max_offset > min_offset, "Encountered an invalid page index buffer"); + } - // Set the first ColumnChunk's offset of ColumnIndex as the adjusted zero offset - int64_t const min_offset = row_groups.front().columns.front().column_index_offset; + file_metadata.setup_page_index(pgidx_bytes, min_offset); + }); +} - // Check if the page index buffer is valid - { - auto const& last_col = row_groups.back().columns.back(); - auto const max_offset = last_col.offset_index_offset + last_col.offset_index_length; - CUDF_EXPECTS(max_offset > min_offset, "Encountered an invalid page index buffer"); +std::vector> aggregate_reader_metadata::all_row_groups( + parquet_reader_options const& options) const +{ + auto const& opts_row_groups = options.get_row_groups(); + if (not opts_row_groups.empty()) { + CUDF_EXPECTS(opts_row_groups.size() == per_file_metadata.size(), + "Row groups in parquet reader options must specify one vector per data source"); + return opts_row_groups; } - file_metadata.setup_page_index(page_index_bytes, min_offset); + std::vector> row_groups; + row_groups.reserve(per_file_metadata.size()); + std::transform(per_file_metadata.begin(), + per_file_metadata.end(), + std::back_inserter(row_groups), + [](auto const& pfm) { + std::vector indices(pfm.row_groups.size()); + std::iota(indices.begin(), indices.end(), size_type{0}); + return indices; + }); + return row_groups; } -size_type aggregate_reader_metadata::total_rows_in_row_groups( +std::size_t aggregate_reader_metadata::total_rows_in_row_groups( cudf::host_span const> row_group_indices) const { - std::size_t total_rows = 0; + CUDF_EXPECTS(row_group_indices.size() == per_file_metadata.size(), + "Encountered unexpected number of input row group indices", + std::invalid_argument); - std::for_each(cuda::counting_iterator{0}, - cuda::counting_iterator{row_group_indices.size()}, - [&](auto const src_idx) { - auto const& pfm = per_file_metadata[src_idx]; - for (auto const row_group_idx : row_group_indices[src_idx]) { - CUDF_EXPECTS(std::cmp_less(row_group_idx, pfm.row_groups.size()), - "Row group index out of bounds"); - total_rows += pfm.row_groups[row_group_idx].num_rows; - } - }); - CUDF_EXPECTS(std::cmp_less_equal(total_rows, std::numeric_limits::max()), - "Total number of rows exceeds cudf::size_type's limit"); - - return static_cast(total_rows); + return std::accumulate( + cuda::counting_iterator{0}, + cuda::counting_iterator{row_group_indices.size()}, + std::size_t{0}, + [&](auto sum, auto const src_idx) { + auto const& file_metadata = per_file_metadata[src_idx]; + return std::accumulate( + row_group_indices[src_idx].begin(), + row_group_indices[src_idx].end(), + sum, + [&](auto sum, auto const row_group_idx) { + CUDF_EXPECTS(std::cmp_greater_equal(row_group_idx, 0) and + std::cmp_less(row_group_idx, file_metadata.row_groups.size()), + std::format("Encountered out-of-bounds row group index for data source. Row " + "group index: {}, Source index: {}, Number of row groups: {}", + row_group_idx, + src_idx, + file_metadata.row_groups.size())); + return sum + file_metadata.row_groups[row_group_idx].num_rows; + }); + }); } std::tuple, @@ -231,8 +287,8 @@ aggregate_reader_metadata::select_payload_columns( return filter_columns_set; }; - // If payload columns are specified, only select payload columns that do not appear in the filter - // expression + // If payload columns are specified, only select payload columns that do not appear in the + // filter expression if (payload_column_names.has_value()) { valid_payload_columns = *payload_column_names; // Remove filter columns from the provided payload column names @@ -634,12 +690,12 @@ std::reference_wrapper named_to_reference_converter::visi { // Map the column index to its name auto const col_name_iter = _column_indices_to_names.find(expr.get_column_index()); - CUDF_EXPECTS( - col_name_iter != _column_indices_to_names.end(), - "Column index in the filter expression not found in the column indices to names map. Note that " - "only top-level columns except structs and lists are supported in " - "Parquet filter expression", - std::invalid_argument); + CUDF_EXPECTS(col_name_iter != _column_indices_to_names.end(), + "Column index in the filter expression not found in the column indices to names " + "map. Note that " + "only top-level columns except structs and lists are supported in " + "Parquet filter expression", + std::invalid_argument); auto const col_name = col_name_iter->second; auto col_index_it = _column_name_to_index.find(col_name); CUDF_EXPECTS(col_index_it != _column_name_to_index.end(), diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_helpers.hpp b/cpp/src/io/parquet/experimental/hybrid_scan_helpers.hpp index df17ae493cd2..e65db678c2d1 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_helpers.hpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_helpers.hpp @@ -80,22 +80,22 @@ class aggregate_reader_metadata : public aggregate_reader_metadata_base { /** * @brief Constructor for aggregate_reader_metadata * - * @param footer_bytes Host span of Parquet file footer buffer bytes + * @param footer_bytes Host span of Parquet file footer buffer bytes, one per source * @param use_arrow_schema Whether to use Arrow schema * @param has_cols_from_mismatched_srcs Whether to have columns from mismatched sources */ - aggregate_reader_metadata(cudf::host_span footer_bytes, + aggregate_reader_metadata(cudf::host_span const> footer_bytes, bool use_arrow_schema, bool has_cols_from_mismatched_srcs); /** * @brief Constructor for aggregate_reader_metadata * - * @param parquet_metadata Pre-populated Parquet file metadata + * @param parquet_metadatas Host span of pre-populated Parquet file metadata, one per source * @param use_arrow_schema Whether to use Arrow schema * @param has_cols_from_mismatched_srcs Whether to have columns from mismatched sources */ - aggregate_reader_metadata(FileMetaData const& parquet_metadata, + aggregate_reader_metadata(cudf::host_span parquet_metadatas, bool use_arrow_schema, bool has_cols_from_mismatched_srcs); @@ -110,21 +110,38 @@ class aggregate_reader_metadata : public aggregate_reader_metadata_base { void initialize_internals(bool use_arrow_schema, bool has_cols_from_mismatched_srcs); /** - * @brief Fetch the byte range of the page index in the Parquet file + * @brief Fetch the byte range of the page index in each Parquet file + * + * @return Vector of byte ranges of the page index, one per source */ - [[nodiscard]] text::byte_range_info page_index_byte_range() const; + [[nodiscard]] std::vector page_index_byte_ranges() const; /** - * @brief Get the Parquet file metadata + * @brief Get the Parquet file metadata for every source + * + * @return Vector of file metadata, one per source */ - [[nodiscard]] FileMetaData parquet_metadata() const; + [[nodiscard]] std::vector parquet_metadatas() const; /** - * @brief Setup and populate the page index structs in `FileMetaData` + * @brief Setup and populate the page index structs in every source's `FileMetaData` + * + * @param page_index_bytes Host span of Parquet page index buffer bytes, one per source + */ + void setup_page_indexes(cudf::host_span const> page_index_bytes); + + /** + * @brief Get all available row group indices, one inner vector per source + * + * If `options.get_row_groups()` is non-empty, validates that its size equals the number of + * sources and returns it as-is. Otherwise returns `[0 .. per_source_num_row_groups[i])` for + * each source. * - * @param page_index_bytes Host span of Parquet page index buffer bytes + * @param options Parquet reader options + * @return Vector of row group indices, one inner vector per source */ - void setup_page_index(cudf::host_span page_index_bytes); + [[nodiscard]] std::vector> all_row_groups( + parquet_reader_options const& options) const; /** * @brief Get the total number of top-level rows in the row groups @@ -132,7 +149,7 @@ class aggregate_reader_metadata : public aggregate_reader_metadata_base { * @param row_group_indices Input row groups indices * @return Total number of top-level rows in the row groups */ - [[nodiscard]] size_type total_rows_in_row_groups( + [[nodiscard]] std::size_t total_rows_in_row_groups( cudf::host_span const> row_group_indices) const; /** diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index 7b53b4345911..f2919c64519f 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -64,10 +64,10 @@ namespace { } // namespace -hybrid_scan_reader_impl::hybrid_scan_reader_impl(cudf::host_span footer_bytes, - parquet_reader_options const& options) +hybrid_scan_reader_impl::hybrid_scan_reader_impl( + cudf::host_span const> footer_bytes, + parquet_reader_options const& options) { - // Open and parse the source dataset metadata _metadata = std::make_unique( footer_bytes, options.is_enabled_use_arrow_schema(), @@ -76,30 +76,30 @@ hybrid_scan_reader_impl::hybrid_scan_reader_impl(cudf::host_span _extended_metadata = static_cast(_metadata.get()); } -hybrid_scan_reader_impl::hybrid_scan_reader_impl(FileMetaData const& parquet_metadata, - parquet_reader_options const& options) +hybrid_scan_reader_impl::hybrid_scan_reader_impl( + cudf::host_span parquet_metadatas, parquet_reader_options const& options) { _metadata = std::make_unique( - parquet_metadata, + parquet_metadatas, options.is_enabled_use_arrow_schema(), options.get_column_names().has_value() and options.is_enabled_allow_mismatched_pq_schemas()); _extended_metadata = static_cast(_metadata.get()); } -FileMetaData hybrid_scan_reader_impl::parquet_metadata() const +std::vector hybrid_scan_reader_impl::parquet_metadatas() const { - return _extended_metadata->parquet_metadata(); + return _extended_metadata->parquet_metadatas(); } -byte_range_info hybrid_scan_reader_impl::page_index_byte_range() const +std::vector hybrid_scan_reader_impl::page_index_byte_ranges() const { - return _extended_metadata->page_index_byte_range(); + return _extended_metadata->page_index_byte_ranges(); } -void hybrid_scan_reader_impl::setup_page_index( - cudf::host_span page_index_bytes) const +void hybrid_scan_reader_impl::setup_page_indexes( + cudf::host_span const> page_index_bytes) const { - _extended_metadata->setup_page_index(page_index_bytes); + _extended_metadata->setup_page_indexes(page_index_bytes); } void hybrid_scan_reader_impl::select_columns(read_columns_mode read_columns_mode, @@ -200,16 +200,13 @@ void hybrid_scan_reader_impl::select_columns(read_columns_mode read_columns_mode [](auto const& buff) { return inline_column_buffer::empty_like(buff); }); } -std::vector hybrid_scan_reader_impl::all_row_groups( +std::vector> hybrid_scan_reader_impl::all_row_groups( parquet_reader_options const& options) const { - auto const num_row_groups = _extended_metadata->get_num_row_groups(); - auto row_groups_indices = std::vector(num_row_groups); - std::iota(row_groups_indices.begin(), row_groups_indices.end(), size_type{0}); - return row_groups_indices; + return _extended_metadata->all_row_groups(options); } -size_type hybrid_scan_reader_impl::total_rows_in_row_groups( +std::size_t hybrid_scan_reader_impl::total_rows_in_row_groups( cudf::host_span const> row_group_indices) const { return _extended_metadata->total_rows_in_row_groups(row_group_indices); diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp index 3f16699a3d18..64eead4463f7 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp @@ -39,49 +39,50 @@ using text::byte_range_info; class hybrid_scan_reader_impl : public parquet::detail::reader_impl { public: /** - * @brief Constructor for the experimental parquet reader implementation to optimally read - * Parquet files subject to highly selective filters + * @brief Constructor for the experimental parquet reader implementation * - * @param footer_bytes Host span of parquet file footer bytes + * @param footer_bytes Span of parquet file footer byte spans, one per source * @param options Parquet reader options */ - explicit hybrid_scan_reader_impl(cudf::host_span footer_bytes, - parquet_reader_options const& options); + explicit hybrid_scan_reader_impl( + cudf::host_span const> footer_bytes, + parquet_reader_options const& options); /** - * @brief Constructor for the experimental parquet reader implementation to optimally read - * Parquet files subject to highly selective filters + * @brief Constructor for the experimental parquet reader implementation * - * @param parquet_metadata Pre-populated Parquet file metadata + * @param parquet_metadatas Span of pre-populated Parquet file metadata, one per source * @param options Parquet reader options */ - explicit hybrid_scan_reader_impl(FileMetaData const& parquet_metadata, + explicit hybrid_scan_reader_impl(cudf::host_span parquet_metadatas, parquet_reader_options const& options); /** - * @copydoc cudf::io::experimental::hybrid_scan::parquet_metadata + * @copydoc cudf::io::experimental::hybrid_scan_multifile::parquet_metadatas */ - [[nodiscard]] FileMetaData parquet_metadata() const; + [[nodiscard]] std::vector parquet_metadatas() const; /** - * @copydoc cudf::io::experimental::hybrid_scan::page_index_byte_range + * @copydoc cudf::io::experimental::hybrid_scan_multifile::page_index_byte_ranges */ - [[nodiscard]] byte_range_info page_index_byte_range() const; + [[nodiscard]] std::vector page_index_byte_ranges() const; /** - * @copydoc cudf::io::experimental::hybrid_scan::setup_page_index + * @copydoc cudf::io::experimental::hybrid_scan_multifile::setup_page_indexes */ - void setup_page_index(cudf::host_span page_index_bytes) const; + void setup_page_indexes( + cudf::host_span const> page_index_bytes) const; /** - * @copydoc cudf::io::experimental::hybrid_scan::all_row_groups + * @copydoc cudf::io::experimental::hybrid_scan_multifile::all_row_groups */ - [[nodiscard]] std::vector all_row_groups(parquet_reader_options const& options) const; + [[nodiscard]] std::vector> all_row_groups( + parquet_reader_options const& options) const; /** * @copydoc cudf::io::experimental::hybrid_scan::total_rows_in_row_groups */ - [[nodiscard]] size_type total_rows_in_row_groups( + [[nodiscard]] std::size_t total_rows_in_row_groups( cudf::host_span const> row_group_indices) const; /** diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp new file mode 100644 index 000000000000..31b3cb5a6443 --- /dev/null +++ b/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp @@ -0,0 +1,60 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "hybrid_scan_impl.hpp" + +#include +#include + +namespace cudf::io::parquet::experimental { + +hybrid_scan_multifile::hybrid_scan_multifile( + cudf::host_span const> footer_bytes, + parquet_reader_options const& options) + : _impl{std::make_unique(footer_bytes, options)} +{ +} + +hybrid_scan_multifile::hybrid_scan_multifile(cudf::host_span parquet_metadata, + parquet_reader_options const& options) + : _impl{std::make_unique(parquet_metadata, options)} +{ +} + +hybrid_scan_multifile::~hybrid_scan_multifile() = default; + +std::vector hybrid_scan_multifile::parquet_metadatas() const +{ + return _impl->parquet_metadatas(); +} + +std::vector hybrid_scan_multifile::page_index_byte_ranges() const +{ + return _impl->page_index_byte_ranges(); +} + +void hybrid_scan_multifile::setup_page_indexes( + cudf::host_span const> page_index_bytes) const +{ + CUDF_FUNC_RANGE(); + _impl->setup_page_indexes(page_index_bytes); +} + +std::vector> hybrid_scan_multifile::all_row_groups( + parquet_reader_options const& options) const +{ + return _impl->all_row_groups(options); +} + +size_type hybrid_scan_multifile::total_rows_in_row_groups( + cudf::host_span const> row_group_indices) const +{ + if (row_group_indices.empty()) { return 0; } + return _impl->total_rows_in_row_groups(row_group_indices); +} + +void hybrid_scan_multifile::reset_column_selection() const { _impl->reset_column_selection(); } + +} // namespace cudf::io::parquet::experimental diff --git a/cpp/src/io/parquet/experimental/page_index_filter.cu b/cpp/src/io/parquet/experimental/page_index_filter.cu index 21ae60443f3d..bb2b89c56c7e 100644 --- a/cpp/src/io/parquet/experimental/page_index_filter.cu +++ b/cpp/src/io/parquet/experimental/page_index_filter.cu @@ -995,12 +995,14 @@ thrust::host_vector aggregate_reader_metadata::compute_data_page_mask( auto const total_rows = total_rows_in_row_groups(row_group_indices); - CUDF_EXPECTS(row_mask_offset + total_rows <= row_mask.size(), - "Mismatch in total rows in input row mask and row groups", - std::invalid_argument); + CUDF_EXPECTS( + std::cmp_less_equal(static_cast(row_mask_offset) + total_rows, row_mask.size()), + "Encountered a mismatch in number of rows in the row group pass and the row mask size", + std::invalid_argument); // Return an empty vector if all rows are invalid or all rows are required - if (row_mask.null_count(row_mask_offset, row_mask_offset + total_rows, stream) == total_rows or + if (std::cmp_equal(row_mask.null_count(row_mask_offset, row_mask_offset + total_rows, stream), + total_rows) or cudf::detail::all_of(row_mask.template begin() + row_mask_offset, row_mask.template begin() + row_mask_offset + total_rows, cuda::std::identity{}, @@ -1091,8 +1093,8 @@ thrust::host_vector aggregate_reader_metadata::compute_data_page_mask( if constexpr (cuda::std::is_same_v) { if (row_mask.nullable() and row_mask.null_count() > 0) { thrust::for_each(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::counting_iterator{row_mask_offset}, - cuda::counting_iterator{row_mask_offset + total_rows}, + cuda::counting_iterator(row_mask_offset), + cuda::counting_iterator(row_mask_offset + total_rows), [row_mask = row_mask.template begin(), null_mask = row_mask.null_mask()] __device__(auto const row_idx) { if (not bit_is_set(null_mask, row_idx)) { row_mask[row_idx] = true; } @@ -1126,7 +1128,7 @@ thrust::host_vector aggregate_reader_metadata::compute_data_page_mask( cudf::detail::make_device_uvector_async(host_tree_level_ptrs, stream, mr); // Build Fenwick tree levels (zeroth level is just the row mask itself) - auto prev_level_size = total_rows; + auto prev_level_size = static_cast(total_rows); std::for_each( cuda::counting_iterator{0}, cuda::counting_iterator{num_levels - 1}, diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index ce7bfdbed0f0..e709a52ace20 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -349,6 +349,7 @@ ConfigureTest( HYBRID_SCAN_TEST io/experimental/hybrid_scan_composer.cpp io/experimental/hybrid_scan_filters_test.cpp + io/experimental/hybrid_scan_multifile_filters_test.cpp io/experimental/hybrid_scan_test.cpp io/parquet_common.cpp io/parquet_test.cpp diff --git a/cpp/tests/io/experimental/hybrid_scan_multifile_filters_test.cpp b/cpp/tests/io/experimental/hybrid_scan_multifile_filters_test.cpp new file mode 100644 index 000000000000..1a0c135207f4 --- /dev/null +++ b/cpp/tests/io/experimental/hybrid_scan_multifile_filters_test.cpp @@ -0,0 +1,220 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "hybrid_scan_common.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace { + +/** + * @brief Struct to hold multifile datasources, and footer buffers along with their byte spans + */ +struct multifile_inputs { + std::vector> datasources; + std::vector> footer_buffers; + std::vector> footer_byte_spans; +}; + +template +multifile_inputs build_multifile_inputs(Buffers const& file_buffers) +{ + multifile_inputs out; + out.datasources.reserve(file_buffers.size()); + out.footer_buffers.reserve(file_buffers.size()); + out.footer_byte_spans.reserve(file_buffers.size()); + for (auto const& buf : file_buffers) { + out.datasources.emplace_back(cudf::io::datasource::create(cudf::host_span( + reinterpret_cast(buf.data()), buf.size()))); + out.footer_buffers.emplace_back( + cudf::io::parquet::fetch_footer_to_host(*out.datasources.back())); + out.footer_byte_spans.emplace_back(*out.footer_buffers.back()); + } + return out; +} + +/** + * @brief Creates a parquet buffer with zero-rows and same schema as table from + * `create_parquet_with_stats` + */ +template +std::vector create_empty_parquet_with_stats() +{ + auto const non_empty = std::get<0>(create_parquet_with_stats()); + auto const empty = cudf::empty_like(non_empty->view()); + + cudf::io::table_input_metadata output_metadata(empty->view()); + output_metadata.column_metadata[0].set_name("col0"); + output_metadata.column_metadata[1].set_name("col1"); + output_metadata.column_metadata[2].set_name("col2"); + + std::vector buffer; + auto out_opts = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&buffer}, empty->view()) + .metadata(std::move(output_metadata)) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .build(); + cudf::io::write_parquet(out_opts); + return buffer; +} + +} // namespace + +struct HybridScanMultifileFiltersTest : public cudf::test::BaseFixture {}; + +TEST_F(HybridScanMultifileFiltersTest, Metadata) +{ + using T = cudf::timestamp_ms; + + // Create two parquet sources, each with 4 row groups and 5000 rows per row + // group + auto constexpr rows_per_row_group = page_size_for_ordered_tests; + auto constexpr num_sources = 2; + + // Build sources with different seeds + std::vector> file_buffers; + file_buffers.reserve(num_sources); + auto constexpr num_concat = 1; + srand(0xbad); + file_buffers.emplace_back(std::get<1>(create_parquet_with_stats())); + srand(0xf00d); + file_buffers.emplace_back(std::get<1>(create_parquet_with_stats())); + + // Filtering AST - col0 < 100 + auto literal_value = + cudf::timestamp_scalar(T(typename T::duration(100)), true, cudf::get_default_stream()); + 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); + + // Construct reader from footer bytes + auto inputs = build_multifile_inputs(file_buffers); + + cudf::io::parquet_reader_options options = + cudf::io::parquet_reader_options::builder().filter(filter_expression).build(); + auto const reader = std::make_unique( + cudf::host_span const>{inputs.footer_byte_spans}, options); + + // Get parquet metadata and check + auto parquet_metadata = reader->parquet_metadatas(); + ASSERT_EQ(parquet_metadata.size(), num_sources); + for (auto const& meta : parquet_metadata) { + ASSERT_FALSE(meta.row_groups.empty()); + EXPECT_FALSE(meta.row_groups[0].columns[0].offset_index.has_value()); + EXPECT_FALSE(meta.row_groups[0].columns[0].column_index.has_value()); + } + + // Setup page index + auto const page_index_byte_ranges = reader->page_index_byte_ranges(); + ASSERT_EQ(page_index_byte_ranges.size(), num_sources); + EXPECT_TRUE(std::all_of(page_index_byte_ranges.begin(), + page_index_byte_ranges.end(), + [](auto const& range) { return not range.is_empty(); })); + + std::vector> page_index_buffers; + std::vector> page_index_byte_spans; + page_index_buffers.reserve(num_sources); + page_index_byte_spans.reserve(num_sources); + + auto iter = cuda::zip_iterator(page_index_byte_ranges.begin(), inputs.datasources.begin()); + std::for_each(iter, iter + num_sources, [&](auto const& pair) { + auto const& [pgidx_byte_range, datasource] = pair; + page_index_buffers.emplace_back( + cudf::io::parquet::fetch_page_index_to_host(*datasource, pgidx_byte_range)); + page_index_byte_spans.emplace_back(*page_index_buffers.back()); + }); + + reader->setup_page_indexes( + cudf::host_span const>{page_index_byte_spans}); + + // Check if page index is now present in each parquet metadata + parquet_metadata = reader->parquet_metadatas(); + for (auto const& meta : parquet_metadata) { + EXPECT_TRUE(meta.row_groups[0].columns[0].offset_index.has_value()); + EXPECT_TRUE(meta.row_groups[0].columns[0].column_index.has_value()); + } + + // Check all row groups + auto input_row_group_indices = reader->all_row_groups(options); + ASSERT_EQ(input_row_group_indices.size(), num_sources); + EXPECT_TRUE(std::all_of( + input_row_group_indices.begin(), input_row_group_indices.end(), [](auto const& rgs) { + return rgs == (std::vector{0, 1, 2, 3}); + })); + + // Set explicit row groups (per-source) via options + options.set_row_groups({{0, 1}, {2, 3}}); + input_row_group_indices = reader->all_row_groups(options); + + // Check if the row groups are set correctly + ASSERT_EQ(input_row_group_indices.size(), num_sources); + EXPECT_EQ(input_row_group_indices[0], (std::vector{0, 1})); + EXPECT_EQ(input_row_group_indices[1], (std::vector{2, 3})); + EXPECT_EQ(reader->total_rows_in_row_groups(input_row_group_indices), + 2 * rows_per_row_group * num_sources); + + // Construct a new reader from a span of existing FileMetaData + auto const reader_with_existing_metadata = + std::make_unique( + cudf::host_span{parquet_metadata}, options); + + // Check if the new metadata is the same as the existing one + auto const new_metadata = reader_with_existing_metadata->parquet_metadatas(); + ASSERT_EQ(new_metadata.size(), num_sources); + EXPECT_TRUE(std::all_of(new_metadata.begin(), new_metadata.end(), [&](auto const& meta) { + return meta.row_groups.size() == parquet_metadata.front().row_groups.size(); + })); +} + +TEST_F(HybridScanMultifileFiltersTest, EmptySource) +{ + using T = uint32_t; + + srand(0xc0ffee); + + // Create two parquet source. First one with non-zero rows and the second one with zero rows. + auto constexpr num_sources = 2; + std::vector> file_buffers; + file_buffers.reserve(num_sources); + file_buffers.emplace_back(std::get<1>(create_parquet_with_stats())); + file_buffers.emplace_back(create_empty_parquet_with_stats()); + + auto inputs = build_multifile_inputs(file_buffers); + + cudf::io::parquet_reader_options options = cudf::io::parquet_reader_options::builder().build(); + auto const reader = std::make_unique( + inputs.footer_byte_spans, options); + + // Check parquet metadata + auto const parquet_metadata = reader->parquet_metadatas(); + ASSERT_EQ(parquet_metadata.size(), num_sources); + EXPECT_FALSE(parquet_metadata.front().row_groups.empty()); + EXPECT_TRUE(parquet_metadata.back().row_groups.empty()); + + // Check row group indices + auto const all_rgs = reader->all_row_groups(options); + ASSERT_EQ(all_rgs.size(), num_sources); + EXPECT_EQ(all_rgs.front(), (std::vector{0, 1, 2, 3})); + EXPECT_TRUE(all_rgs.back().empty()); + + // Check page index byte ranges + auto const page_index_byte_ranges = reader->page_index_byte_ranges(); + ASSERT_EQ(page_index_byte_ranges.size(), num_sources); + EXPECT_FALSE(page_index_byte_ranges.front().is_empty()); + EXPECT_TRUE(page_index_byte_ranges.back().is_empty()); +}