diff --git a/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp b/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp index c75fa3d186d..23e4b5b4412 100644 --- a/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp +++ b/cpp/include/cudf/io/experimental/hybrid_scan_multifile.hpp @@ -61,6 +61,8 @@ class hybrid_scan_multifile { /** * @brief Constructor for the multi-file experimental Parquet reader * + * @throws std::invalid_argument if no sources are provided + * * @param footer_bytes Host span of Parquet file footer byte spans, one per source * @param options Parquet reader options */ @@ -70,12 +72,25 @@ class hybrid_scan_multifile { /** * @brief Constructor for the multi-file experimental Parquet reader * + * @throws std::invalid_argument if no sources are provided + * * @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 Constructor that takes ownership of pre-populated Parquet file metadata + * + * @throws std::invalid_argument if no sources are provided + * + * @param parquet_metadata Pre-populated Parquet file metadata, one per source + * @param options Parquet reader options + */ + explicit hybrid_scan_multifile(std::vector&& parquet_metadata, + parquet_reader_options const& options); + /** * @brief Destructor for the multi-file experimental Parquet reader */ diff --git a/cpp/src/io/parquet/experimental/hybrid_scan.cpp b/cpp/src/io/parquet/experimental/hybrid_scan.cpp index c133ed3130e..bd6a159b6ea 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan.cpp @@ -18,7 +18,7 @@ hybrid_scan_metadata::hybrid_scan_metadata(cudf::host_span footer : _metadata{std::make_shared( std::vector>{footer_bytes}, options.is_enabled_use_arrow_schema(), - options.get_column_names().has_value() and options.is_enabled_allow_mismatched_pq_schemas())} + options.is_enabled_allow_mismatched_pq_schemas())} { } @@ -27,7 +27,7 @@ hybrid_scan_metadata::hybrid_scan_metadata(FileMetaData const& parquet_metadata, : _metadata{std::make_shared( std::vector{parquet_metadata}, options.is_enabled_use_arrow_schema(), - options.get_column_names().has_value() and options.is_enabled_allow_mismatched_pq_schemas())} + options.is_enabled_allow_mismatched_pq_schemas())} { } diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_helpers.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_helpers.cpp index faba65f4f33..d0523b55fe7 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_helpers.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_helpers.cpp @@ -36,22 +36,6 @@ using text::byte_range_info; namespace { -// Construct a vector of FileMetaData from the input footer bytes -[[nodiscard]] std::vector parquet_metadatas_from_footer_bytes( - cudf::host_span const> footer_bytes) -{ - std::vector parquet_metadatas; - parquet_metadatas.reserve(footer_bytes.size()); - std::transform(footer_bytes.begin(), - footer_bytes.end(), - std::back_inserter(parquet_metadatas), - [](auto const& footer_bytes) { - metadata parsed_metadata{footer_bytes}; - return FileMetaData{std::move(parsed_metadata)}; - }); - return parquet_metadatas; -} - // Construct a vector of all row group indices from the input vectors [[nodiscard]] auto all_row_group_indices( std::span const> row_group_indices) @@ -132,25 +116,31 @@ 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(parquet_metadatas_from_footer_bytes(footer_bytes), - use_arrow_schema, - has_cols_from_mismatched_srcs) + : aggregate_reader_metadata( + parquet::detail::parallel_construct_metadatas( + footer_bytes, [](auto const& bytes) { return FileMetaData{metadata{bytes}}; }), + use_arrow_schema, + has_cols_from_mismatched_srcs) { - CUDF_EXPECTS( - not footer_bytes.empty(), "At least one source must be provided", std::invalid_argument); } 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( + : aggregate_reader_metadata( std::vector{parquet_metadatas.begin(), parquet_metadatas.end()}, use_arrow_schema, has_cols_from_mismatched_srcs) { - CUDF_EXPECTS( - not parquet_metadatas.empty(), "At least one source must be provided", std::invalid_argument); +} + +aggregate_reader_metadata::aggregate_reader_metadata(std::vector&& parquet_metadatas, + bool use_arrow_schema, + bool has_cols_from_mismatched_srcs) + : aggregate_reader_metadata_base( + std::move(parquet_metadatas), use_arrow_schema, has_cols_from_mismatched_srcs) +{ } std::vector aggregate_reader_metadata::page_index_byte_ranges() const diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_helpers.hpp b/cpp/src/io/parquet/experimental/hybrid_scan_helpers.hpp index e0a3746ecf5..075d85660d9 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_helpers.hpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_helpers.hpp @@ -81,6 +81,8 @@ class aggregate_reader_metadata : public aggregate_reader_metadata_base { /** * @brief Constructor for aggregate_reader_metadata * + * @throws std::invalid_argument if no sources are provided + * * @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 @@ -92,6 +94,8 @@ class aggregate_reader_metadata : public aggregate_reader_metadata_base { /** * @brief Constructor for aggregate_reader_metadata * + * @throws std::invalid_argument if no sources are provided + * * @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 @@ -100,6 +104,19 @@ class aggregate_reader_metadata : public aggregate_reader_metadata_base { bool use_arrow_schema, bool has_cols_from_mismatched_srcs); + /** + * @brief Constructor that takes ownership of pre-populated Parquet file metadata + * + * @throws std::invalid_argument if no sources are provided + * + * @param parquet_metadatas 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(std::vector&& parquet_metadatas, + bool use_arrow_schema, + bool has_cols_from_mismatched_srcs); + aggregate_reader_metadata(aggregate_reader_metadata const&) = delete; aggregate_reader_metadata& operator=(aggregate_reader_metadata const&) = delete; aggregate_reader_metadata(aggregate_reader_metadata&&) = default; diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp index 61e159f9f83..59cd164d3e6 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp @@ -180,8 +180,10 @@ hybrid_scan_reader_impl::hybrid_scan_reader_impl( cudf::host_span const> footer_bytes, parquet_reader_options const& options) { - _metadata = std::make_shared( - footer_bytes, options.is_enabled_use_arrow_schema(), has_cols_from_mismatched_sources(options)); + _metadata = + std::make_shared(footer_bytes, + options.is_enabled_use_arrow_schema(), + options.is_enabled_allow_mismatched_pq_schemas()); _extended_metadata = static_cast(_metadata.get()); } @@ -192,7 +194,17 @@ hybrid_scan_reader_impl::hybrid_scan_reader_impl( _metadata = std::make_shared(parquet_metadatas, options.is_enabled_use_arrow_schema(), - has_cols_from_mismatched_sources(options)); + options.is_enabled_allow_mismatched_pq_schemas()); + _extended_metadata = static_cast(_metadata.get()); +} + +hybrid_scan_reader_impl::hybrid_scan_reader_impl(std::vector&& parquet_metadatas, + parquet_reader_options const& options) +{ + _metadata = + std::make_shared(std::move(parquet_metadatas), + options.is_enabled_use_arrow_schema(), + options.is_enabled_allow_mismatched_pq_schemas()); _extended_metadata = static_cast(_metadata.get()); } @@ -285,12 +297,13 @@ void hybrid_scan_reader_impl::select_columns(read_columns_mode read_columns_mode // Save original output-buffer schema for reuse across materialization passes. _original_output_buffers_template = make_empty_like_column_buffers(_output_buffers); - // Initialize mutable output-buffer template for this materialization pass. - reset_output_buffers_template(); + // Initialize mutable output buffers for this materialization pass. + reset_output_buffers(); } -void hybrid_scan_reader_impl::reset_output_buffers_template() +void hybrid_scan_reader_impl::reset_output_buffers() { + _output_buffers = make_empty_like_column_buffers(_original_output_buffers_template); _output_buffers_template = make_empty_like_column_buffers(_original_output_buffers_template); } @@ -336,7 +349,7 @@ void hybrid_scan_reader_impl::prepare_materialization(read_columns_mode read_col reset_internal_state(); initialize_options(options, num_sources, stream, mr); select_columns(read_columns_mode, options); - reset_output_buffers_template(); + reset_output_buffers(); } std::vector> diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp b/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp index bbdacd818a4..c3feb436b82 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_impl.hpp @@ -58,6 +58,15 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { explicit hybrid_scan_reader_impl(cudf::host_span parquet_metadatas, parquet_reader_options const& options); + /** + * @brief Constructor that takes ownership of pre-populated Parquet file metadata + * + * @param parquet_metadatas Pre-populated Parquet file metadata, one per source + * @param options Parquet reader options + */ + explicit hybrid_scan_reader_impl(std::vector&& parquet_metadatas, + parquet_reader_options const& options); + /** * @brief Constructor that takes shared ownership of pre-parsed Parquet metadata * @@ -393,9 +402,9 @@ class hybrid_scan_reader_impl : public parquet::detail::reader_impl { void mark_buffers_nullable_for_pruned_pages(); /** - * @brief Initialize the mutable output-buffer template for this materialization + * @brief Reset the output buffers and their template from the original selected-columns schema */ - void reset_output_buffers_template(); + void reset_output_buffers(); /** * @brief Select the columns to be read based on the read mode diff --git a/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp b/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp index 259435401aa..493c86a8f3b 100644 --- a/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp +++ b/cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp @@ -26,6 +26,12 @@ hybrid_scan_multifile::hybrid_scan_multifile(cudf::host_span { } +hybrid_scan_multifile::hybrid_scan_multifile(std::vector&& parquet_metadata, + parquet_reader_options const& options) + : _impl{std::make_unique(std::move(parquet_metadata), options)} +{ +} + hybrid_scan_multifile::~hybrid_scan_multifile() = default; std::vector hybrid_scan_multifile::parquet_metadatas() const diff --git a/cpp/src/io/parquet/reader_impl.hpp b/cpp/src/io/parquet/reader_impl.hpp index 6fbd4789c6d..137d43c0316 100644 --- a/cpp/src/io/parquet/reader_impl.hpp +++ b/cpp/src/io/parquet/reader_impl.hpp @@ -413,19 +413,6 @@ class reader_impl { _file_itm_data._current_input_pass < _file_itm_data.num_passes(); } - /** - * @brief Check if the user has specified columns from mismatched sources - * - * @param options Reader options - * @return True if the user has specified columns from mismatched sources - */ - [[nodiscard]] bool has_cols_from_mismatched_sources(parquet_reader_options const& options) const - { - return (options.get_column_names().has_value() or - options.get_column_field_ids().has_value()) and - options.is_enabled_allow_mismatched_pq_schemas(); - } - /** * @brief Effective `ignore_missing_columns` policy for column selection * @@ -440,6 +427,20 @@ class reader_impl { not(has_cols_from_mismatched_sources(options) and _metadata->get_num_sources() > 1); } + private: + /** + * @brief Check if the user has specified columns from mismatched sources + * + * @param options Reader options + * @return True if the user has specified columns from mismatched sources + */ + [[nodiscard]] bool has_cols_from_mismatched_sources(parquet_reader_options const& options) const + { + return (options.get_column_names().has_value() or + options.get_column_field_ids().has_value()) and + options.is_enabled_allow_mismatched_pq_schemas(); + } + protected: /** * @brief Check if the user has specified custom row bounds diff --git a/cpp/src/io/parquet/reader_impl_chunking.cu b/cpp/src/io/parquet/reader_impl_chunking.cu index 1cb95bd35a9..f728ebfc00c 100644 --- a/cpp/src/io/parquet/reader_impl_chunking.cu +++ b/cpp/src/io/parquet/reader_impl_chunking.cu @@ -19,6 +19,7 @@ #include #include +#include namespace cudf::io::parquet::detail { @@ -419,29 +420,32 @@ void reader_impl::create_global_chunk_info() auto const num_chunks = row_groups_info.size() * num_input_columns; // Mapping of input column to page index column - std::vector column_mapping; - - if (_has_offset_index and not row_groups_info.empty()) { - // use first row group to define mappings (assumes same schema for each file) - auto const& rg = row_groups_info[0]; - auto const& columns = _metadata->get_row_group(rg.index, rg.source_index).columns; - column_mapping.resize(num_input_columns); - std::transform( - _input_columns.begin(), _input_columns.end(), column_mapping.begin(), [&](auto const& col) { - // translate schema_idx into something we can use for the page indexes - if (auto it = std::find_if(columns.begin(), - columns.end(), - [&](auto const& col_chunk) { - return col_chunk.schema_idx == - _metadata->map_schema_index(col.schema_idx, - rg.source_index); - }); - it != columns.end()) { - return std::distance(columns.begin(), it); - } - CUDF_FAIL("cannot find column mapping"); - }); - } + auto column_mappings = std::unordered_map>{}; + + auto const column_mapping_for_source = [&](auto const& rg) -> std::vector const& { + auto const [iter, inserted] = column_mappings.try_emplace(rg.source_index); + if (inserted) { + auto const& columns = _metadata->get_row_group(rg.index, rg.source_index).columns; + auto& mapping = iter->second; + mapping.resize(num_input_columns); + std::transform( + _input_columns.begin(), _input_columns.end(), mapping.begin(), [&](auto const& col) { + // translate schema_idx into something we can use for the page indexes + if (auto it = std::find_if(columns.begin(), + columns.end(), + [&](auto const& col_chunk) { + return col_chunk.schema_idx == + _metadata->map_schema_index(col.schema_idx, + rg.source_index); + }); + it != columns.end()) { + return static_cast(std::distance(columns.begin(), it)); + } + CUDF_FAIL("cannot find column mapping"); + }); + } + return iter->second; + }; // Initialize column chunk information auto remaining_rows = num_rows; @@ -454,6 +458,8 @@ void reader_impl::create_global_chunk_info() auto row_group_rows = std::min(remaining_rows + adjusted_row_group_rows, row_group.num_rows); + auto const* const column_mapping = _has_offset_index ? &column_mapping_for_source(rg) : nullptr; + // generate ColumnChunkDesc objects for everything to be decoded (all input columns) for (size_t i = 0; i < num_input_columns; ++i) { auto col = _input_columns[i]; @@ -479,7 +485,7 @@ void reader_impl::create_global_chunk_info() // grab the column_chunk_info for each chunk (if it exists) column_chunk_info const* const chunk_info = - _has_offset_index ? &rg.column_chunks.value()[column_mapping[i]] : nullptr; + _has_offset_index ? &rg.column_chunks.value()[(*column_mapping)[i]] : nullptr; chunks.emplace_back(col_meta.total_compressed_size, nullptr, diff --git a/cpp/src/io/parquet/reader_impl_helpers.cpp b/cpp/src/io/parquet/reader_impl_helpers.cpp index f229a612072..ad5d1236aa6 100644 --- a/cpp/src/io/parquet/reader_impl_helpers.cpp +++ b/cpp/src/io/parquet/reader_impl_helpers.cpp @@ -290,6 +290,20 @@ struct schema_child_lookup { std::unordered_map schema_idx_caches; }; +/** + * @brief Marks a field in the destination schema as optional if it is optional in the source schema + * + * @param dst Destination schema element + * @param src Source schema element + */ +void propagate_optional_field(SchemaElement& dst, SchemaElement const& src) +{ + if (dst.repetition_type == FieldRepetitionType::REQUIRED and + src.repetition_type != FieldRepetitionType::REQUIRED) { + dst.repetition_type = FieldRepetitionType::OPTIONAL; + } +} + } // namespace type_id to_type_id(SchemaElement const& schema, @@ -510,7 +524,17 @@ void metadata::sanitize_schema() process(0); } -metadata::metadata(FileMetaData&& other) : FileMetaData(std::move(other)) {} +metadata::metadata(FileMetaData&& other) : FileMetaData(std::move(other)) +{ + // Since page index is set up for all or no row groups, just check if any column chunk has it set. + // Update this check if this behavior changes in the future. + is_page_index_setup = + std::any_of(row_groups.cbegin(), row_groups.cend(), [](auto const& row_group) { + return std::any_of(row_group.columns.cbegin(), row_group.columns.cend(), [](auto const& col) { + return col.column_index.has_value() or col.offset_index.has_value(); + }); + }); +} metadata::metadata(datasource* source, bool read_page_indexes) { @@ -550,6 +574,8 @@ metadata::metadata(datasource* source, bool read_page_indexes) void metadata::setup_page_index(cudf::host_span page_index_bytes, int64_t min_offset) { + if (is_page_index_setup) { return; } + CUDF_FUNC_RANGE(); // Flatten all columns into a single vector for easier task distribution @@ -625,6 +651,8 @@ void metadata::setup_page_index(cudf::host_span page_index_bytes, read_column_indexes(cp, col_ref.get()); } } + + is_page_index_setup = true; } metadata::~metadata() @@ -645,26 +673,9 @@ metadata::~metadata() std::vector aggregate_reader_metadata::metadatas_from_sources( host_span const> sources, bool read_page_indexes) { - // Avoid using the thread pool for a single source - if (sources.size() == 1) { - std::vector result; - result.emplace_back(sources[0].get(), read_page_indexes); - return result; - } - - std::vector> metadata_ctor_tasks; - metadata_ctor_tasks.reserve(sources.size()); - for (auto const& source : sources) { - metadata_ctor_tasks.emplace_back(cudf::detail::host_worker_pool().submit_task( - [source = source.get(), read_page_indexes] { return metadata{source, read_page_indexes}; })); - } - std::vector metadatas; - metadatas.reserve(sources.size()); - std::transform(metadata_ctor_tasks.begin(), - metadata_ctor_tasks.end(), - std::back_inserter(metadatas), - [](std::future& task) { return std::move(task).get(); }); - return metadatas; + return parallel_construct_metadatas(sources, [read_page_indexes](auto const& source) { + return metadata{source.get(), read_page_indexes}; + }); } std::vector> @@ -926,23 +937,19 @@ void aggregate_reader_metadata::initialize_internals(bool use_arrow_schema, } CUDF_EXPECTS(schema == pfm.schema, "All sources must have the same schema"); } - } - // Mark the column schema in the first (default) source as nullable if it is nullable in any of - // the input sources. This avoids recomputing this within build_column() and - // populate_metadata(). - std::for_each( - cuda::counting_iterator{static_cast(1)}, - cuda::counting_iterator{schema.size()}, - [&](auto const schema_idx) { - if (schema[schema_idx].repetition_type == FieldRepetitionType::REQUIRED and - std::any_of( - per_file_metadata.begin() + 1, per_file_metadata.end(), [&](auto const& pfm) { - return pfm.schema[schema_idx].repetition_type != FieldRepetitionType::REQUIRED; - })) { - schema[schema_idx].repetition_type = FieldRepetitionType::OPTIONAL; - } - }); + // Mark a field in the first source's schema as nullable if it is nullable in any other + // source + std::for_each( + cuda::counting_iterator{static_cast(1)}, + cuda::counting_iterator{schema.size()}, + [&](auto const schema_idx) { + std::for_each( + per_file_metadata.begin() + 1, per_file_metadata.end(), [&](auto const& pfm) { + propagate_optional_field(schema[schema_idx], pfm.schema[schema_idx]); + }); + }); + } } // Collect and apply arrow:schema from Parquet's key value metadata section @@ -957,6 +964,10 @@ aggregate_reader_metadata::aggregate_reader_metadata(std::vector&& bool use_arrow_schema, bool has_cols_from_mismatched_srcs) { + CUDF_EXPECTS(not parquet_metadatas.empty(), + "Encountered an empty vector of parquet metadatas (sources)", + std::invalid_argument); + per_file_metadata.reserve(parquet_metadatas.size()); std::transform(std::make_move_iterator(parquet_metadatas.begin()), std::make_move_iterator(parquet_metadatas.end()), @@ -982,6 +993,10 @@ aggregate_reader_metadata::aggregate_reader_metadata( num_rows(calc_num_rows()), num_row_groups(calc_num_row_groups()) { + CUDF_EXPECTS(not per_file_metadata.empty(), + "Encountered an empty vector of parquet sources", + std::invalid_argument); + initialize_internals(use_arrow_schema, has_cols_from_mismatched_srcs); } @@ -1967,6 +1982,8 @@ aggregate_reader_metadata::select_columns( auto const case_sensitive_names = selection_options.case_sensitive_names; auto const selection_mode = selection_options.selection_mode; + auto constexpr root_idx = 0; + // Setup schema lookup helper auto schema_lookup = schema_child_lookup{[&](int const schema_idx, int const src_idx) -> SchemaElement const& { @@ -2087,14 +2104,14 @@ aggregate_reader_metadata::select_columns( }; // Compares two schema elements to be equal except their number of children - auto const equal_to_except_num_children = [selection_mode](SchemaElement const& lhs, - SchemaElement const& rhs) { + auto const equal_to_except_num_children = [selection_mode, case_sensitive_names]( + SchemaElement const& lhs, SchemaElement const& rhs) { // Match by field ID if enabled, otherwise match by name auto const match_schema_by_field_id = selection_mode == column_selection_mode::BY_FIELD_ID; auto const names_match = (match_schema_by_field_id and lhs.field_id.has_value() and rhs.field_id.has_value()) ? lhs.field_id == rhs.field_id - : lhs.name == rhs.name; + : are_column_paths_equal(lhs.name, rhs.name, case_sensitive_names); return lhs.type == rhs.type and lhs.converted_type == rhs.converted_type and lhs.type_length == rhs.type_length and names_match and lhs.decimal_scale == rhs.decimal_scale and @@ -2127,6 +2144,9 @@ aggregate_reader_metadata::select_columns( // Map the schema index from 0th tree (src) to the one in the current (dst) tree. schema_idx_map[src_schema_idx] = dst_schema_idx; + // Mark the field nullable in the first schema if it is nullable in the current tree. + propagate_optional_field(per_file_metadata.front().schema[src_schema_idx], dst_schema_elem); + // If src_schema_elem is a stub, it does not exist in the column_name_info and column_buffer // hierarchy. So continue on with mapping. if (src_schema_elem.is_stub()) { @@ -2191,6 +2211,29 @@ aggregate_reader_metadata::select_columns( } }; + // Maps a top-level column's schema_idx across the rest of the data sources if we are reading from + // mismatched Parquet sources. `col_name_info` is null when all of the column's children are + // selected. + auto map_column_across_sources = [&](column_name_info const* col_name_info, + std::string const& col_name, + int const src_schema_idx) { + if (per_file_metadata.size() == 1 or schema_idx_maps.empty()) { return; } + + std::for_each( + cuda::counting_iterator{static_cast(1)}, + cuda::counting_iterator{per_file_metadata.size()}, + [&](auto const src_idx) { + // Ensure that each top level column exists in the destination schema tree. + auto const dst_schema_idx = + schema_lookup.find_target_schema_child(root_idx, root_idx, col_name, src_idx); + CUDF_EXPECTS( + dst_schema_idx != -1, + std::format("Encountered missing top-level column '{}' across Parquet sources", col_name), + std::invalid_argument); + map_column(col_name_info, src_schema_idx, dst_schema_idx, src_idx); + }); + }; + std::vector output_column_schemas; // @@ -2215,6 +2258,7 @@ aggregate_reader_metadata::select_columns( auto const& root = get_schema(0); if (not use_names.has_value()) { for (auto const& schema_idx : root.children_idx) { + map_column_across_sources(nullptr, get_schema(schema_idx).name, schema_idx); build_column(nullptr, schema_idx, output_columns, false); output_column_schemas.push_back(schema_idx); } @@ -2330,31 +2374,14 @@ aggregate_reader_metadata::select_columns( } } } + + // Map the column's schema_idx across the rest of the data sources and propagate nullability. for (auto& col : selected_columns) { - auto constexpr root_idx = 0; - auto const& top_level_col_schema_idx = + auto const top_level_col_schema_idx = schema_lookup.find_schema_child_by_name(root_idx, col.name); - bool const valid_column = build_column(&col, top_level_col_schema_idx, output_columns, false); - if (valid_column) { + map_column_across_sources(&col, col.name, top_level_col_schema_idx); + if (build_column(&col, top_level_col_schema_idx, output_columns, false)) { output_column_schemas.push_back(top_level_col_schema_idx); - - // Map the column's schema_idx across the rest of the data sources if required. - if (per_file_metadata.size() > 1 and not schema_idx_maps.empty()) { - std::for_each( - cuda::counting_iterator{static_cast(1)}, - cuda::counting_iterator{per_file_metadata.size()}, - [&](auto const src_idx) { - // Ensure that each top level column exists in the destination schema tree. - auto const dst_col_schema_idx = - schema_lookup.find_target_schema_child(root_idx, root_idx, col.name, src_idx); - CUDF_EXPECTS( - dst_col_schema_idx != -1, - std::format("Encountered missing top-level column '{}' across Parquet sources", - col.name), - std::invalid_argument); - map_column(&col, top_level_col_schema_idx, dst_col_schema_idx, src_idx); - }); - } } } } diff --git a/cpp/src/io/parquet/reader_impl_helpers.hpp b/cpp/src/io/parquet/reader_impl_helpers.hpp index 50279594164..b8f804ee6b3 100644 --- a/cpp/src/io/parquet/reader_impl_helpers.hpp +++ b/cpp/src/io/parquet/reader_impl_helpers.hpp @@ -7,26 +7,79 @@ #include "parquet_gpu.hpp" +#include #include #include #include #include #include +#include +#include #include +#include #include +#include #include #include #include #include #include #include +#include #include #include #include namespace cudf::io::parquet::detail { +/** + * @brief Construct metadatas from inputs using the host worker pool for multiple inputs + * + * @param inputs Metadata construction inputs, one per source + * @param op Operation constructing a metadata object from one input + * @return Constructed metadata objects, in input order + */ +template +[[nodiscard]] auto parallel_construct_metadatas(cudf::host_span inputs, UnaryOp op) +{ + using result_type = std::invoke_result_t; + + std::vector results; + results.reserve(inputs.size()); + + // Avoid using the thread pool for a single input + if (inputs.size() == 1) { + results.emplace_back(op(inputs.front())); + return results; + } + + std::vector> tasks; + tasks.reserve(inputs.size()); + + auto pending_exception = std::exception_ptr{}; + try { + std::transform(inputs.begin(), inputs.end(), std::back_inserter(tasks), [&op](T const& input) { + return cudf::detail::host_worker_pool().submit_task( + [&op, input_ptr = &input] { return op(*input_ptr); }); + }); + } catch (...) { + pending_exception = std::current_exception(); + } + + for (auto& task : tasks) { + try { + results.emplace_back(task.get()); + } catch (...) { + if (not pending_exception) { pending_exception = std::current_exception(); } + } + } + + if (pending_exception) { std::rethrow_exception(pending_exception); } + + return results; +} + /** * @brief page location and size info */ @@ -147,6 +200,9 @@ struct metadata : public FileMetaData { protected: void sanitize_schema(); + + private: + bool is_page_index_setup = false; }; /** diff --git a/cpp/tests/io/experimental/hybrid_scan_multifile_filters_test.cpp b/cpp/tests/io/experimental/hybrid_scan_multifile_filters_test.cpp index 6e7cdad91bd..fbf28c85bb0 100644 --- a/cpp/tests/io/experimental/hybrid_scan_multifile_filters_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_multifile_filters_test.cpp @@ -180,16 +180,17 @@ TEST_F(HybridScanMultifileFiltersTest, Metadata) 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 + // Move the existing FileMetaData into a new reader without copying it. + auto const num_row_groups = parquet_metadata.front().row_groups.size(); auto const reader_with_existing_metadata = std::make_unique( - cudf::host_span{parquet_metadata}, options); + std::move(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(); + return meta.row_groups.size() == num_row_groups; })); } diff --git a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp index 5d3eb9e8083..804c754a157 100644 --- a/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_multifile_test.cpp @@ -501,3 +501,100 @@ TEST_F(HybridScanMultifileTest, AllColumnsPreservesRequiredNullability) EXPECT_FALSE(result.tbl->view().column(0).nullable()); CUDF_TEST_EXPECT_TABLES_EQUAL(input_table->view(), result.tbl->view()); } + +TEST_F(HybridScanMultifileTest, ReadColumnsFromMismatchedSchemas) +{ + // Create two sources with mismatched schemas + auto const buffer_a = std::get<1>(create_parquet_with_stats()); + auto const buffer_b = std::get<1>(create_parquet_with_stats( + 100, cudf::io::compression_type::AUTO, {"col2", "col0", "col1"}, {2, 0, 1})); + + auto const parquet_buffers = std::vector>{buffer_a, buffer_b}; + auto const source_info = build_source_info(parquet_buffers); + auto inputs = multifile_inputs(source_info); + auto const stream = cudf::get_default_stream(); + auto const mr = cudf::get_current_device_resource_ref(); + + // Reading mismatched schemas must be opted into, even without a column projection + EXPECT_THROW(cudf::io::parquet::experimental::hybrid_scan_multifile( + inputs.footer_byte_spans, cudf::io::parquet_reader_options::builder().build()), + cudf::logic_error); + + // Expected table from the regular reader + auto const expected = + cudf::io::read_parquet(cudf::io::parquet_reader_options::builder(source_info) + .allow_mismatched_pq_schemas(true) + .column_names({"col0", "col1", "col2"}) + .build(), + stream, + mr); + + auto options = + cudf::io::parquet_reader_options::builder().allow_mismatched_pq_schemas(true).build(); + auto const reader = + cudf::io::parquet::experimental::hybrid_scan_multifile{inputs.footer_byte_spans, options}; + auto const row_groups = reader.all_row_groups(options); + + // Single step materialize with hybrid scan + { + auto column_data = fetch_multisource_device_data( + inputs, reader.all_column_chunks_byte_ranges(row_groups, options), stream, mr); + auto const result = + reader.materialize_all_columns(row_groups, column_data.flat_spans, options, stream, mr); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected.tbl->view(), result.tbl->view()); + } + + // Two step materialize with hybrid scan + { + auto literal_value = cudf::numeric_scalar(std::numeric_limits::min()); + auto literal = cudf::ast::literal(literal_value); + auto col_ref = cudf::ast::column_name_reference("col0"); + auto filter = cudf::ast::operation(cudf::ast::ast_operator::GREATER_EQUAL, col_ref, literal); + + options.set_filter(filter); + reader.reset_column_selection(); + + auto row_mask = reader.build_all_true_row_mask(row_groups, stream, mr); + auto row_mask_view = row_mask->mutable_view(); + + auto filter_column_chunks = fetch_multisource_device_data( + inputs, reader.filter_column_chunks_byte_ranges(row_groups, options), stream, mr); + auto const filter_result = reader.materialize_filter_columns(row_groups, + filter_column_chunks.flat_spans, + row_mask_view, + use_data_page_mask::NO, + options, + stream, + mr); + + auto payload_column_chunks = fetch_multisource_device_data( + inputs, reader.payload_column_chunks_byte_ranges(row_groups, options), stream, mr); + auto const payload_result = reader.materialize_payload_columns(row_groups, + payload_column_chunks.flat_spans, + row_mask_view, + use_data_page_mask::NO, + options, + stream, + mr); + + CUDF_TEST_EXPECT_TABLES_EQUAL(expected.tbl->select({0}), filter_result.tbl->view()); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected.tbl->select({1, 2}), payload_result.tbl->view()); + } +} + +TEST_F(HybridScanMultifileTest, EmptySources) +{ + // Arrow schema is applied during metadata construction, so make sure empty inputs are rejected + // before any metadata is touched + auto const options = cudf::io::parquet_reader_options::builder().use_arrow_schema(true).build(); + + EXPECT_THROW(cudf::io::parquet::experimental::hybrid_scan_multifile( + cudf::host_span const>{}, options), + std::invalid_argument); + EXPECT_THROW(cudf::io::parquet::experimental::hybrid_scan_multifile( + cudf::host_span{}, options), + std::invalid_argument); + EXPECT_THROW(cudf::io::parquet::experimental::hybrid_scan_multifile( + std::vector{}, options), + std::invalid_argument); +} diff --git a/cpp/tests/io/parquet_reader_test.cpp b/cpp/tests/io/parquet_reader_test.cpp index e1d47c1fe91..470d2dab07e 100644 --- a/cpp/tests/io/parquet_reader_test.cpp +++ b/cpp/tests/io/parquet_reader_test.cpp @@ -690,23 +690,24 @@ TEST_F(ParquetReaderTest, SelectMismatchedStructChildByFieldId) .metadata(std::move(metadata_a)); cudf::io::write_parquet(write_args_a); - auto y_b = cudf::test::fixed_width_column_wrapper{40, 50}; + auto y_b = cudf::test::fixed_width_column_wrapper{{40, 50}, {true, false}}; auto x_b = cudf::test::fixed_width_column_wrapper{4, 5}; auto struct_b = cudf::test::structs_column_wrapper{{y_b, x_b}, {true, true}}.release(); cudf::table_view const table_b{{*struct_b}}; auto path_b = temp_env->get_temp_filepath("SelectNestedFieldIdChildOrderB.parquet"); cudf::io::table_input_metadata metadata_b(table_b); - metadata_b.column_metadata[0].set_name("record").set_parquet_field_id(1); - metadata_b.column_metadata[0].child(0).set_name("y").set_parquet_field_id(3); - metadata_b.column_metadata[0].child(1).set_name("x").set_parquet_field_id(2); + metadata_b.column_metadata[0].set_name("renamed_record").set_parquet_field_id(1); + metadata_b.column_metadata[0].child(0).set_name("renamed_y").set_parquet_field_id(3); + metadata_b.column_metadata[0].child(1).set_name("renamed_x").set_parquet_field_id(2); auto write_args_b = cudf::io::parquet_writer_options::builder(cudf::io::sink_info{path_b}, table_b) .metadata(std::move(metadata_b)); cudf::io::write_parquet(write_args_b); auto expected_x = cudf::test::fixed_width_column_wrapper{1, 2, 3, 4, 5}; - auto expected_y = cudf::test::fixed_width_column_wrapper{10, 20, 30, 40, 50}; + auto expected_y = cudf::test::fixed_width_column_wrapper{ + {10, 20, 30, 40, 50}, {true, true, true, true, false}}; auto expected_struct = cudf::test::structs_column_wrapper{{expected_x, expected_y}, {true, true, true, true, true}} .release(); @@ -5492,6 +5493,17 @@ TEST_F(ParquetReaderTest, LateBindSourceInfo) CUDF_TEST_EXPECT_TABLES_EQUAL(result.tbl->view(), expected->view()); } +TEST_F(ParquetReaderTest, EmptySourcesWithArrowSchema) +{ + auto sources = std::vector>{}; + auto file_metadatas = std::vector{}; + auto const options = cudf::io::parquet_reader_options::builder(cudf::io::source_info{}) + .use_arrow_schema(true) + .build(); + EXPECT_THROW(cudf::io::read_parquet(std::move(sources), std::move(file_metadatas), options), + std::invalid_argument); +} + TEST_F(ParquetReaderTest, InvalidFooterMagic) { auto const expected = create_random_fixed_table(4, 4, false); @@ -6372,6 +6384,30 @@ TEST_F(ParquetReaderTest, NestedMismatchedSchemaColumnValidation) EXPECT_THROW(cudf::io::read_parquet(opts), std::invalid_argument); } } + +TEST_F(ParquetReaderTest, CaseInsensitiveMismatchedSchemasPropagateNullability) +{ + auto const required = cudf::test::fixed_width_column_wrapper{1, 2, 3}; + auto const optional = cudf::test::fixed_width_column_wrapper{{4, 5}, {true, false}}; + auto const required_path = + write_parquet_temp_file(cudf::table_view{{required}}, "CaseRequired.parquet", {"column"}); + auto const optional_path = + write_parquet_temp_file(cudf::table_view{{optional}}, "CaseOptional.parquet", {"COLUMN"}); + + auto const options = + cudf::io::parquet_reader_options::builder(cudf::io::source_info{{required_path, optional_path}}) + .allow_mismatched_pq_schemas(true) + .case_sensitive_names(false) + .column_names({"column"}) + .build(); + + // A non-nullable column in the first source but nullable in another must be read as nullable. + auto const expected = cudf::test::fixed_width_column_wrapper{ + {1, 2, 3, 4, 5}, {true, true, true, true, false}}; + auto const result = cudf::io::read_parquet(options); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result.tbl->view().column(0), expected); +} + namespace { /**