diff --git a/cpp/include/cudf/io/parquet_io_utils.hpp b/cpp/include/cudf/io/parquet_io_utils.hpp index 97e2fe97b5db..ba557f1f8ca6 100644 --- a/cpp/include/cudf/io/parquet_io_utils.hpp +++ b/cpp/include/cudf/io/parquet_io_utils.hpp @@ -12,6 +12,7 @@ #include #include +#include #include #include #include @@ -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 * diff --git a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp index 7170d818cd87..61ca6c7f4038 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include #include @@ -22,12 +23,18 @@ #include #include +#include +#include +#include +#include #include #include #include #include +#include #include #include +#include /** * @file parquet_io_utils.cpp @@ -86,23 +93,58 @@ auto dispatch_fetch_tasks(std::size_t num_sources, Task fetch_task) std::vector> fetch_footers_to_host_impl( cudf::host_span 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(header_buffer->data()); - auto ender_buffer = datasource.host_read(len - ender_len, ender_len); - auto const ender = reinterpret_cast(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(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( + speculative_buffer->data() + speculative_buffer->size() - ender_len); + + if (speculative_read_offset == 0) { + auto const header = reinterpret_cast(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 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 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("LIBCUDF_PARQUET_METADATA_SIZE_HINT", + default_metadata_size_hint); +} + std::unique_ptr 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, 1> datasources{std::ref(datasource)}; auto footer_buffers = fetch_footers_to_host_impl({datasources.data(), datasources.size()}); return std::move(footer_buffers.front()); diff --git a/cpp/tests/io/parquet_reader_test.cpp b/cpp/tests/io/parquet_reader_test.cpp index 8ca9de9c10db..32a16027026a 100644 --- a/cpp/tests/io/parquet_reader_test.cpp +++ b/cpp/tests/io/parquet_reader_test.cpp @@ -28,10 +28,13 @@ #include #include +#include #include +#include #include #include #include +#include using ParquetDecompressionTest = DecompressionTest; @@ -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(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(ender_buffer->data()); + auto const footer_len = static_cast(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 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;