Improve parquet footer reading performance by speculatively reading footer bytes

cpp/include/cudf/io/parquet_io_utils.hpp

@@ -12,6 +12,7 @@
 #include <rmm/device_buffer.hpp>
 #include <rmm/resource_ref.hpp>
 
+#include <cstddef>
 #include <functional>
 #include <future>
 #include <tuple>
@@ -34,6 +35,25 @@ namespace io::parquet {
 //! Using `byte_range_info` from cudf::io::text
 using cudf::io::text::byte_range_info;
 
+/**
+ * @brief Returns the Parquet reader's footer speculative read size in bytes.
+ *
+ * @ingroup io_utils
+ *
+ * Controlled by the `LIBCUDF_PARQUET_METADATA_SIZE_HINT` environment variable.
+ * Defaults to 64 KiB.
+ *
+ * When the footer is smaller than the speculative read size, the footer metadata
+ * is loaded in a single read, which is especially useful for high-latency, remote
+ * storage systems. When the footer is larger than the speculative read size, the
+ * footer metadata will be loaded in two reads.
+ *
+ * Set `LIBCUDF_PARQUET_METADATA_SIZE_HINT=0` to disable speculative reads.
+ *
+ * @return Number of bytes to speculatively read from the end of the source.
+ */
+[[nodiscard]] std::size_t metadata_size_hint();
+
 /**
  * @brief Fetches a host buffer of Parquet footer bytes from the input data source
  *

cpp/src/io/parquet/io_utils/parquet_io_utils.cpp

@@ -8,6 +8,7 @@
 
 #include <cudf/detail/nvtx/ranges.hpp>
 #include <cudf/detail/utilities/cuda_memcpy.hpp>
+#include <cudf/detail/utilities/getenv_or.hpp>
 #include <cudf/detail/utilities/host_worker_pool.hpp>
 #include <cudf/detail/utilities/integer_utils.hpp>
 #include <cudf/io/datasource.hpp>
@@ -22,12 +23,18 @@
 #include <cuda/iterator>
 #include <cuda/std/tuple>
 
+#include <algorithm>
+#include <array>
+#include <cstring>
+#include <format>
 #include <functional>
 #include <mutex>
 #include <numeric>
 #include <stdexcept>
+#include <string>
 #include <tuple>
 #include <type_traits>
+#include <vector>
 
 /**
  * @file parquet_io_utils.cpp
@@ -86,23 +93,58 @@ auto dispatch_fetch_tasks(std::size_t num_sources, Task fetch_task)
 std::vector<std::unique_ptr<cudf::io::datasource::buffer>> fetch_footers_to_host_impl(
   cudf::host_span<std::reference_wrapper<cudf::io::datasource> const> datasources)
 {
+  // Look up runtime configuration once, as late as possible.
+  auto const metadata_size_hint = cudf::io::parquet::metadata_size_hint();
   // Helper to fetch footer from a datasource
-  auto const fetch_footer = [](cudf::io::datasource& datasource) {
+  auto const fetch_footer = [metadata_size_hint](cudf::io::datasource& datasource) {
     constexpr auto header_len = sizeof(file_header_s);
     constexpr auto ender_len  = sizeof(file_ender_s);
     size_t const len          = datasource.size();
     CUDF_EXPECTS(len > header_len + ender_len, "Incorrect data source");
 
-    auto header_buffer = datasource.host_read(0, header_len);
-    auto const header  = reinterpret_cast<file_header_s const*>(header_buffer->data());
-    auto ender_buffer  = datasource.host_read(len - ender_len, ender_len);
-    auto const ender   = reinterpret_cast<file_ender_s const*>(ender_buffer->data());
-    CUDF_EXPECTS(header->magic == detail::parquet_magic, "Corrupted header");
+    auto const speculative_read_size =
+      std::min(len, std::max(metadata_size_hint, static_cast<size_t>(ender_len)));
+    auto const speculative_read_offset = len - speculative_read_size;
+
+    auto speculative_buffer = datasource.host_read(speculative_read_offset, speculative_read_size);
+    CUDF_EXPECTS(speculative_buffer->size() == speculative_read_size,
+                 "Failed to read Parquet speculative metadata bytes");
+
+    auto const ender = reinterpret_cast<file_ender_s const*>(
+      speculative_buffer->data() + speculative_buffer->size() - ender_len);
+
+    if (speculative_read_offset == 0) {
+      auto const header = reinterpret_cast<file_header_s const*>(speculative_buffer->data());
+      CUDF_EXPECTS(header->magic == detail::parquet_magic, "Corrupted header");
+    }
+
     CUDF_EXPECTS(ender->magic == detail::parquet_magic, "Corrupted footer");
     CUDF_EXPECTS(ender->footer_len != 0 && ender->footer_len <= (len - header_len - ender_len),
                  "Incorrect footer length");
 
-    return datasource.host_read(len - ender->footer_len - ender_len, ender->footer_len);
+    auto const footer_offset = len - ender->footer_len - ender_len;
+    if (footer_offset >= speculative_read_offset) {
+      // fastpath: the speculative read includes the full footer.
+      auto const footer_start_offset = footer_offset - speculative_read_offset;
+      CUDF_EXPECTS(footer_start_offset + ender->footer_len <= speculative_buffer->size(),
+                   "Speculative metadata read did not include full footer bytes");
+      std::vector<uint8_t> footer_bytes(ender->footer_len);
+      std::memcpy(
+        footer_bytes.data(), speculative_buffer->data() + footer_start_offset, ender->footer_len);
+      return cudf::io::datasource::buffer::create(std::move(footer_bytes));
+    }
+
+    // The speculative read only got part of the footer. Read the missing prefix, then stitch.
+    auto const missing_prefix_size = speculative_read_offset - footer_offset;
+    auto missing_prefix            = datasource.host_read(footer_offset, missing_prefix_size);
+    CUDF_EXPECTS(missing_prefix->size() == missing_prefix_size,
+                 "Failed to read the missing footer prefix bytes");
+    std::vector<uint8_t> footer_bytes(ender->footer_len);
+    std::memcpy(footer_bytes.data(), missing_prefix->data(), missing_prefix_size);
+    auto const footer_suffix_size = ender->footer_len - missing_prefix_size;
+    std::memcpy(
+      footer_bytes.data() + missing_prefix_size, speculative_buffer->data(), footer_suffix_size);
+    return cudf::io::datasource::buffer::create(std::move(footer_bytes));
   };
 
   return dispatch_fetch_tasks(datasources.size(), [&](std::size_t source_idx) {
@@ -335,11 +377,16 @@ fetch_byte_ranges_to_device_async_impl(
 
 }  // namespace
 
+[[nodiscard]] std::size_t metadata_size_hint()
+{
+  static constexpr auto default_metadata_size_hint = std::size_t{64} * 1024;
+  return cudf::detail::getenv_or<std::size_t>("LIBCUDF_PARQUET_METADATA_SIZE_HINT",
+                                              default_metadata_size_hint);
+}
+
 std::unique_ptr<cudf::io::datasource::buffer> fetch_footer_to_host(cudf::io::datasource& datasource)
 {
   CUDF_FUNC_RANGE();
-
-  // Wrap the input into an array and delegate to the multi-source implementation
   std::array<std::reference_wrapper<cudf::io::datasource>, 1> datasources{std::ref(datasource)};
   auto footer_buffers = fetch_footers_to_host_impl({datasources.data(), datasources.size()});
   return std::move(footer_buffers.front());

cpp/tests/io/parquet_reader_test.cpp

@@ -28,10 +28,13 @@
 #include <src/io/parquet/parquet_gpu.hpp>
 #include <src/io/parquet/stats_filter_helpers.hpp>
 
+#include <algorithm>
 #include <array>
+#include <cstring>
 #include <limits>
 #include <memory>
 #include <stdexcept>
+#include <utility>
 
 using ParquetDecompressionTest = DecompressionTest<ParquetReaderTest>;
 
@@ -2872,6 +2875,44 @@ struct ParquetMetadataReaderTest : public cudf::test::BaseFixture {
   }
 };
 
+static constexpr char const* parquet_metadata_size_hint_env_var =
+  "LIBCUDF_PARQUET_METADATA_SIZE_HINT";
+
+struct ParquetMetadataSizeHintTest : public ParquetReaderTest {};
+
+TEST_F(ParquetMetadataSizeHintTest, ReadParquet)
+{
+  srand(31337);
+  auto const expected = create_random_fixed_table<int>(2, 8, false);
+
+  auto const filepath = temp_env->get_temp_filepath("MetadataSizeHint.parquet");
+  cudf::io::parquet_writer_options write_opts =
+    cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, *expected);
+  cudf::io::write_parquet(write_opts);
+
+  auto const source        = cudf::io::datasource::create(filepath);
+  auto const file_size     = source->size();
+  auto constexpr ender_len = sizeof(cudf::io::parquet::file_ender_s);
+  auto const ender_buffer  = source->host_read(file_size - ender_len, ender_len);
+  auto const ender = reinterpret_cast<cudf::io::parquet::file_ender_s const*>(ender_buffer->data());
+  auto const footer_len = static_cast<size_t>(ender->footer_len);
+
+  auto const source_info = cudf::io::source_info{filepath};
+
+  // Test cases:
+  // - 0: disable speculative reads
+  // - 9: smaller than typical footer metadata
+  // - footer_len + 1: larger than the footer
+  // - file_size + 1: larger than the entire file
+  std::vector<size_t> const hints{0, 9, footer_len + 1, file_size + 1};
+  for (auto const hint : hints) {
+    tmp_env_var const env(parquet_metadata_size_hint_env_var, std::to_string(hint));
+    auto const read_opts = cudf::io::parquet_reader_options::builder(source_info).build();
+    auto const result    = cudf::io::read_parquet(read_opts);
+    CUDF_TEST_EXPECT_TABLES_EQUAL(result.tbl->view(), expected->view());
+  }
+}
+
 TEST_F(ParquetMetadataReaderTest, Basics)
 {
   auto const num_rows = 1200;