Hybrid scan constructor that moves in materialized footers

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<FileMetaData const> 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<FileMetaData>&& parquet_metadata,
+                                 parquet_reader_options const& options);
+
   /**
    * @brief Destructor for the multi-file experimental Parquet reader
    */

cpp/src/io/parquet/experimental/hybrid_scan.cpp

@@ -18,7 +18,7 @@ hybrid_scan_metadata::hybrid_scan_metadata(cudf::host_span<uint8_t const> footer
   : _metadata{std::make_shared<detail::aggregate_reader_metadata>(
       std::vector<cudf::host_span<uint8_t const>>{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<detail::aggregate_reader_metadata>(
       std::vector<FileMetaData>{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())}
 {
 }
 

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<FileMetaData> parquet_metadatas_from_footer_bytes(
-  cudf::host_span<cudf::host_span<uint8_t const> const> footer_bytes)
-{
-  std::vector<FileMetaData> 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<std::vector<cudf::size_type> const> row_group_indices)
@@ -132,25 +116,31 @@ aggregate_reader_metadata::aggregate_reader_metadata(
   cudf::host_span<cudf::host_span<uint8_t const> 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<FileMetaData const> parquet_metadatas,
   bool use_arrow_schema,
   bool has_cols_from_mismatched_srcs)
-  : aggregate_reader_metadata_base(
+  : aggregate_reader_metadata(
       std::vector<FileMetaData>{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<FileMetaData>&& 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<text::byte_range_info> aggregate_reader_metadata::page_index_byte_ranges() const

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<FileMetaData>&& 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;

cpp/src/io/parquet/experimental/hybrid_scan_impl.cpp

@@ -180,8 +180,10 @@ hybrid_scan_reader_impl::hybrid_scan_reader_impl(
   cudf::host_span<cudf::host_span<uint8_t const> const> footer_bytes,
   parquet_reader_options const& options)
 {
-  _metadata = std::make_shared<aggregate_reader_metadata>(
-    footer_bytes, options.is_enabled_use_arrow_schema(), has_cols_from_mismatched_sources(options));
+  _metadata =
+    std::make_shared<aggregate_reader_metadata>(footer_bytes,
+                                                options.is_enabled_use_arrow_schema(),
+                                                options.is_enabled_allow_mismatched_pq_schemas());
 
   _extended_metadata = static_cast<aggregate_reader_metadata*>(_metadata.get());
 }
@@ -192,7 +194,17 @@ hybrid_scan_reader_impl::hybrid_scan_reader_impl(
   _metadata =
     std::make_shared<aggregate_reader_metadata>(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<aggregate_reader_metadata*>(_metadata.get());
+}
+
+hybrid_scan_reader_impl::hybrid_scan_reader_impl(std::vector<FileMetaData>&& parquet_metadatas,
+                                                 parquet_reader_options const& options)
+{
+  _metadata =
+    std::make_shared<aggregate_reader_metadata>(std::move(parquet_metadatas),
+                                                options.is_enabled_use_arrow_schema(),
+                                                options.is_enabled_allow_mismatched_pq_schemas());
   _extended_metadata = static_cast<aggregate_reader_metadata*>(_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<std::vector<cudf::size_type>>

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<FileMetaData const> 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<FileMetaData>&& parquet_metadatas,
+                                   parquet_reader_options const& options);
+
   /**
    * @brief Constructor that takes shared ownership of pre-parsed Parquet metadata
    *
@@ -400,9 +409,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

cpp/src/io/parquet/experimental/hybrid_scan_multifile.cpp

@@ -26,6 +26,12 @@ hybrid_scan_multifile::hybrid_scan_multifile(cudf::host_span<FileMetaData const>
 {
 }
 
+hybrid_scan_multifile::hybrid_scan_multifile(std::vector<FileMetaData>&& parquet_metadata,
+                                             parquet_reader_options const& options)
+  : _impl{std::make_unique<detail::hybrid_scan_reader_impl>(std::move(parquet_metadata), options)}
+{
+}
+
 hybrid_scan_multifile::~hybrid_scan_multifile() = default;
 
 std::vector<FileMetaData> hybrid_scan_multifile::parquet_metadatas() const

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

cpp/src/io/parquet/reader_impl_chunking.cu

@@ -19,6 +19,7 @@
 #include <thrust/transform_scan.h>
 
 #include <numeric>
+#include <unordered_map>
 
 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<size_type> 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<size_type, std::vector<size_type>>{};
+
+  auto const column_mapping_for_source = [&](auto const& rg) -> std::vector<size_type> 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<size_type>(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<size_t>(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,