Add python bindings for hybrid scan metadata and release GIL in the reader APIs

cpp/include/cudf/io/experimental/hybrid_scan.hpp

@@ -32,6 +32,11 @@ namespace cudf::io::parquet::experimental::detail {
  *        Hybrid Scan operation.
  */
 class hybrid_scan_reader_impl;
+
+/**
+ * @brief Internal parsed Parquet file metadata for the Hybrid Scan reader.
+ */
+class aggregate_reader_metadata;
 }  // namespace cudf::io::parquet::experimental::detail
 
 //! Using `byte_range_info` from cudf::io::text
@@ -53,6 +58,70 @@ enum class use_data_page_mask : bool {
   NO  = false  ///< Do not compute or use a data page mask
 };
 
+/**
+ * @brief Shareable, pre-parsed Parquet file metadata for the Hybrid Scan reader.
+ *
+ * Parses the Parquet file metadata once so that multiple `hybrid_scan_reader` instances reading
+ * the same file can share it rather than each re-parsing and copying the row group metadata.
+ * The intended use is to read disjoint row-group ranges of a single file: construct one
+ * `hybrid_scan_metadata` per file and pass it to as many readers as there are ranges.
+ *
+ * @code{.cpp}
+ * // Parse the metadata once
+ * auto metadata = parquet::experimental::hybrid_scan_metadata{*footer_buffer, options};
+ * // Construct lightweight readers that share it
+ * auto reader_a = std::make_unique<parquet::experimental::hybrid_scan_reader>(metadata);
+ * auto reader_b = std::make_unique<parquet::experimental::hybrid_scan_reader>(metadata);
+ * @endcode
+ *
+ * @note The metadata is immutable after `setup_page_index()` has been called (or immediately after
+ * construction if page index setup is skipped). Concurrent usage by multiple readers is thread
+ * safe. This handle does not support multi-source (multi-file) metadata.
+ */
+class hybrid_scan_metadata {
+ public:
+  /**
+   * @brief Parse and own Parquet file metadata from a span of footer bytes
+   *
+   * @param footer_bytes Host span of Parquet file footer bytes
+   * @param options Parquet reader options
+   */
+  hybrid_scan_metadata(cudf::host_span<uint8_t const> footer_bytes,
+                       parquet_reader_options const& options);
+
+  /**
+   * @brief Own Parquet file metadata from a pre-populated `FileMetaData`
+   *
+   * @param parquet_metadata Pre-populated Parquet file metadata
+   * @param options Parquet reader options
+   */
+  hybrid_scan_metadata(FileMetaData const& parquet_metadata, parquet_reader_options const& options);
+
+  /**
+   * @brief Destructor for the shared Parquet metadata
+   */
+  ~hybrid_scan_metadata();
+
+  hybrid_scan_metadata(hybrid_scan_metadata const&) = default;  ///< Copy constructor
+  hybrid_scan_metadata(hybrid_scan_metadata&&)      = default;  ///< Move constructor
+
+  /**
+   * @brief Copy assignment operator
+   * @return Reference to this object
+   */
+  hybrid_scan_metadata& operator=(hybrid_scan_metadata const&) = default;
+
+  /**
+   * @brief Move assignment operator
+   * @return Reference to this object
+   */
+  hybrid_scan_metadata& operator=(hybrid_scan_metadata&&) = default;
+
+ private:
+  std::shared_ptr<detail::aggregate_reader_metadata> _metadata;
+  friend class hybrid_scan_reader;
+};
+
 /**
  * @brief The experimental parquet reader class to optimally read parquet files subject to
  *        highly selective filters, called a Hybrid Scan operation
@@ -301,6 +370,15 @@ class hybrid_scan_reader {
   explicit hybrid_scan_reader(FileMetaData const& parquet_metadata,
                               parquet_reader_options const& options);
 
+  /**
+   * @brief Constructor that takes shared ownership of pre-parsed Parquet file metadata
+   *
+   * Constructs a reader that shares the pre-parsed metadata object.
+   *
+   * @param metadata Shared, pre-parsed Parquet file metadata
+   */
+  explicit hybrid_scan_reader(hybrid_scan_metadata metadata);
+
   /**
    * @brief Destructor for the experimental parquet reader class
    */

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

@@ -13,6 +13,26 @@
 
 namespace cudf::io::parquet::experimental {
 
+hybrid_scan_metadata::hybrid_scan_metadata(cudf::host_span<uint8_t const> footer_bytes,
+                                           parquet_reader_options const& options)
+  : _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())}
+{
+}
+
+hybrid_scan_metadata::hybrid_scan_metadata(FileMetaData const& parquet_metadata,
+                                           parquet_reader_options const& options)
+  : _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())}
+{
+}
+
+hybrid_scan_metadata::~hybrid_scan_metadata() = default;
+
 hybrid_scan_reader::hybrid_scan_reader(cudf::host_span<uint8_t const> footer_bytes,
                                        parquet_reader_options const& options)
   : _impl{std::make_unique<detail::hybrid_scan_reader_impl>(
@@ -27,6 +47,11 @@ hybrid_scan_reader::hybrid_scan_reader(FileMetaData const& parquet_metadata,
 {
 }
 
+hybrid_scan_reader::hybrid_scan_reader(hybrid_scan_metadata metadata)
+  : _impl{std::make_unique<detail::hybrid_scan_reader_impl>(std::move(metadata._metadata))}
+{
+}
+
 hybrid_scan_reader::~hybrid_scan_reader() = default;
 
 [[nodiscard]] text::byte_range_info hybrid_scan_reader::page_index_byte_range() const

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

@@ -122,7 +122,7 @@ 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_unique<aggregate_reader_metadata>(
+  _metadata = std::make_shared<aggregate_reader_metadata>(
     footer_bytes, options.is_enabled_use_arrow_schema(), has_cols_from_mismatched_sources(options));
 
   _extended_metadata = static_cast<aggregate_reader_metadata*>(_metadata.get());
@@ -132,12 +132,20 @@ hybrid_scan_reader_impl::hybrid_scan_reader_impl(
   cudf::host_span<FileMetaData const> parquet_metadatas, parquet_reader_options const& options)
 {
   _metadata =
-    std::make_unique<aggregate_reader_metadata>(parquet_metadatas,
+    std::make_shared<aggregate_reader_metadata>(parquet_metadatas,
                                                 options.is_enabled_use_arrow_schema(),
                                                 has_cols_from_mismatched_sources(options));
   _extended_metadata = static_cast<aggregate_reader_metadata*>(_metadata.get());
 }
 
+hybrid_scan_reader_impl::hybrid_scan_reader_impl(
+  std::shared_ptr<aggregate_reader_metadata> metadata)
+{
+  CUDF_EXPECTS(metadata != nullptr, "Shared parquet metadata must not be null");
+  _metadata          = std::move(metadata);
+  _extended_metadata = static_cast<aggregate_reader_metadata*>(_metadata.get());
+}
+
 std::vector<FileMetaData> hybrid_scan_reader_impl::parquet_metadatas() const
 {
   return _extended_metadata->parquet_metadatas();

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

@@ -58,6 +58,13 @@ 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 shared ownership of pre-parsed Parquet metadata
+   *
+   * @param metadata Shared, pre-parsed Parquet file metadata. Must not be null.
+   */
+  explicit hybrid_scan_reader_impl(std::shared_ptr<aggregate_reader_metadata> metadata);
+
   /**
    * @copydoc cudf::io::parquet::experimental::hybrid_scan_multifile::parquet_metadatas
    */

cpp/src/io/parquet/reader_impl.cpp

@@ -550,11 +550,11 @@ reader_impl::reader_impl(std::size_t chunk_read_limit,
   CUDF_EXPECTS(file_metadatas.empty() or file_metadatas.size() == _sources.size(),
                "Encountered a mismatch in the number of provided data sources and metadatas");
 
-  _metadata = file_metadatas.empty() ? std::make_unique<aggregate_reader_metadata>(
+  _metadata = file_metadatas.empty() ? std::make_shared<aggregate_reader_metadata>(
                                          _sources,
                                          options.is_enabled_use_arrow_schema(),
                                          has_cols_from_mismatched_sources(options))
-                                     : std::make_unique<aggregate_reader_metadata>(
+                                     : std::make_shared<aggregate_reader_metadata>(
                                          std::forward<std::vector<FileMetaData>>(file_metadatas),
                                          options.is_enabled_use_arrow_schema(),
                                          has_cols_from_mismatched_sources(options));

cpp/src/io/parquet/reader_impl.hpp

@@ -571,7 +571,7 @@ class reader_impl {
   named_to_reference_converter _expr_conv{std::nullopt, table_metadata{}, true};
 
   std::vector<std::unique_ptr<datasource>> _sources;
-  std::unique_ptr<aggregate_reader_metadata> _metadata;
+  std::shared_ptr<aggregate_reader_metadata> _metadata;
 
   // Number of sources
   size_t _num_sources{0};

cpp/tests/io/experimental/hybrid_scan_common.cpp

@@ -420,6 +420,7 @@ std::pair<std::unique_ptr<cudf::table>, std::vector<char>> create_parquet_with_s
   INSTANTIATE_CREATE_PARQUET_WITH_STATS(T, 1, true, true)
 
 INSTANTIATE_CREATE_PARQUET_WITH_STATS(uint32_t, 4, true, false);
+INSTANTIATE_CREATE_PARQUET_WITH_STATS(int32_t, 2, true, false);
 INSTANTIATE_CREATE_PARQUET_WITH_STATS(cudf::timestamp_ms, 2, true, false);
 INSTANTIATE_CREATE_PARQUET_WITH_STATS(cudf::duration_ms, 2, true, false);
 

cpp/tests/io/experimental/hybrid_scan_test.cpp

@@ -985,6 +985,175 @@ TEST_F(HybridScanTest, StructChildFilterColumn)
     std::invalid_argument);
 }
 
+TEST_F(HybridScanTest, SharedMetadataReaderMatchesReadParquet)
+{
+  using T                              = int32_t;
+  auto constexpr num_concat            = 2;
+  auto [written_table, parquet_buffer] = create_parquet_with_stats<T, num_concat>();
+
+  auto const stream  = cudf::get_default_stream();
+  auto const mr      = cudf::get_current_device_resource_ref();
+  auto const options = cudf::io::parquet_reader_options::builder().build();
+
+  auto datasource          = cudf::io::datasource::create(cudf::host_span<std::byte const>(
+    reinterpret_cast<std::byte const*>(parquet_buffer.data()), parquet_buffer.size()));
+  auto const footer_buffer = cudf::io::parquet::fetch_footer_to_host(*datasource);
+
+  // Parse the file metadata once and share it across independent readers.
+  auto const metadata =
+    cudf::io::parquet::experimental::hybrid_scan_metadata{*footer_buffer, options};
+
+  // Read all columns (single step) through a reader that borrows the shared metadata.
+  auto const read_all_columns = [&] {
+    auto const reader =
+      std::make_unique<cudf::io::parquet::experimental::hybrid_scan_reader>(metadata);
+    auto const row_groups   = reader->all_row_groups(options);
+    auto const chunk_ranges = reader->all_column_chunks_byte_ranges(row_groups, options);
+    auto [buffers, data, tasks] =
+      cudf::io::parquet::fetch_byte_ranges_to_device_async(*datasource, chunk_ranges, stream, mr);
+    tasks.get();
+    return reader->materialize_all_columns(row_groups, data, options, stream, mr).tbl;
+  };
+
+  // Two readers sharing one metadata instance each produce the same table as the main reader.
+  auto const table_a = read_all_columns();
+  auto const table_b = read_all_columns();
+
+  auto const expected =
+    cudf::io::read_parquet(
+      cudf::io::parquet_reader_options::builder(
+        cudf::io::source_info(cudf::host_span<char>(parquet_buffer.data(), parquet_buffer.size())))
+        .build(),
+      stream)
+      .tbl;
+
+  CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), table_a->view());
+  CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), table_b->view());
+}
+
+TEST_F(HybridScanTest, SharedMetadataFromFileMetaDataMatchesReadParquet)
+{
+  using T                              = int32_t;
+  auto constexpr num_concat            = 2;
+  auto [written_table, parquet_buffer] = create_parquet_with_stats<T, num_concat>();
+
+  auto const stream  = cudf::get_default_stream();
+  auto const mr      = cudf::get_current_device_resource_ref();
+  auto const options = cudf::io::parquet_reader_options::builder().build();
+
+  auto datasource = cudf::io::datasource::create(cudf::host_span<std::byte const>(
+    reinterpret_cast<std::byte const*>(parquet_buffer.data()), parquet_buffer.size()));
+
+  // Obtain FileMetaData from an initial reader, then build shared metadata from it.
+  auto const footer_buffer = cudf::io::parquet::fetch_footer_to_host(*datasource);
+  auto const seed_reader =
+    std::make_unique<cudf::io::parquet::experimental::hybrid_scan_reader>(*footer_buffer, options);
+  auto const file_metadata = seed_reader->parquet_metadata();
+
+  auto const metadata =
+    cudf::io::parquet::experimental::hybrid_scan_metadata{file_metadata, options};
+
+  // Two readers sharing the FileMetaData-derived metadata each produce the correct table.
+  auto const read_all_columns = [&] {
+    auto const reader =
+      std::make_unique<cudf::io::parquet::experimental::hybrid_scan_reader>(metadata);
+    auto const row_groups   = reader->all_row_groups(options);
+    auto const chunk_ranges = reader->all_column_chunks_byte_ranges(row_groups, options);
+    auto [buffers, data, tasks] =
+      cudf::io::parquet::fetch_byte_ranges_to_device_async(*datasource, chunk_ranges, stream, mr);
+    tasks.get();
+    return reader->materialize_all_columns(row_groups, data, options, stream, mr).tbl;
+  };
+
+  auto const table_a = read_all_columns();
+  auto const table_b = read_all_columns();
+
+  auto const expected =
+    cudf::io::read_parquet(
+      cudf::io::parquet_reader_options::builder(
+        cudf::io::source_info(cudf::host_span<char>(parquet_buffer.data(), parquet_buffer.size())))
+        .build(),
+      stream)
+      .tbl;
+
+  CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), table_a->view());
+  CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), table_b->view());
+}
+
+TEST_F(HybridScanTest, SharedMetadataConcurrentReadersMatchReadParquet)
+{
+  using T                              = uint32_t;
+  auto constexpr num_concat            = 4;
+  auto [written_table, parquet_buffer] = create_parquet_with_stats<T, num_concat>();
+
+  auto const stream  = cudf::get_default_stream();
+  auto const mr      = cudf::get_current_device_resource_ref();
+  auto const options = cudf::io::parquet_reader_options::builder().build();
+
+  auto datasource          = cudf::io::datasource::create(cudf::host_span<std::byte const>(
+    reinterpret_cast<std::byte const*>(parquet_buffer.data()), parquet_buffer.size()));
+  auto const footer_buffer = cudf::io::parquet::fetch_footer_to_host(*datasource);
+
+  // Parse the metadata once and share it, by value, across two readers that are alive at the
+  // same time, each responsible for a disjoint range of the file's row groups.
+  auto const metadata =
+    cudf::io::parquet::experimental::hybrid_scan_metadata{*footer_buffer, options};
+  auto const reader_a =
+    std::make_unique<cudf::io::parquet::experimental::hybrid_scan_reader>(metadata);
+  auto const reader_b =
+    std::make_unique<cudf::io::parquet::experimental::hybrid_scan_reader>(metadata);
+
+  // Only one of the readers sharing this metadata sets up the page index.
+  auto const page_index_byte_range = reader_a->page_index_byte_range();
+  ASSERT_FALSE(page_index_byte_range.is_empty());
+  auto const page_index_buffer =
+    cudf::io::parquet::fetch_page_index_to_host(*datasource, page_index_byte_range);
+  reader_a->setup_page_index(*page_index_buffer);
+
+  // The page index materialized through `reader_a` must be visible through `reader_b` since both
+  // readers share the same underlying metadata.
+  auto const metadata_from_b = reader_b->parquet_metadata();
+  ASSERT_GT(metadata_from_b.row_groups.size(), 1);
+  for (auto const& row_group : metadata_from_b.row_groups) {
+    for (auto const& column_chunk : row_group.columns) {
+      EXPECT_TRUE(column_chunk.column_index.has_value());
+      EXPECT_TRUE(column_chunk.offset_index.has_value());
+    }
+  }
+
+  // Split the row groups into two disjoint ranges and read each range through a different reader
+  // sharing the metadata, with both readers alive and used concurrently.
+  auto const all_row_groups = reader_a->all_row_groups(options);
+  auto const split          = all_row_groups.size() / 2;
+  auto const row_groups_a =
+    std::vector<cudf::size_type>(all_row_groups.begin(), all_row_groups.begin() + split);
+  auto const row_groups_b =
+    std::vector<cudf::size_type>(all_row_groups.begin() + split, all_row_groups.end());
+
+  auto const materialize = [&](auto const& reader, auto const& row_group_indices) {
+    auto const chunk_ranges = reader->all_column_chunks_byte_ranges(row_group_indices, options);
+    auto [buffers, data, tasks] =
+      cudf::io::parquet::fetch_byte_ranges_to_device_async(*datasource, chunk_ranges, stream, mr);
+    tasks.get();
+    return reader->materialize_all_columns(row_group_indices, data, options, stream, mr).tbl;
+  };
+
+  auto const table_a = materialize(reader_a, row_groups_a);
+  auto const table_b = materialize(reader_b, row_groups_b);
+
+  auto const table =
+    cudf::concatenate(std::vector<cudf::table_view>{table_a->view(), table_b->view()});
+  auto const expected =
+    cudf::io::read_parquet(
+      cudf::io::parquet_reader_options::builder(
+        cudf::io::source_info(cudf::host_span<char>(parquet_buffer.data(), parquet_buffer.size())))
+        .build(),
+      stream)
+      .tbl;
+
+  CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), table->view());
+}
+
 TEST_F(HybridScanTest, AllRowsPrunedReportsInputRowGroups)
 {
   using cudf::io::parquet::experimental::use_data_page_mask;