From 1bd7284f83d5ebf6b921a72de0cb9822e0aa55c4 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 9 Jun 2026 13:17:39 -0700 Subject: [PATCH 1/4] Merge tom/libcudf-speculative-footer-read --- cpp/include/cudf/io/config_utils.hpp | 26 ++- cpp/include/cudf/io/detail/parquet.hpp | 1 + .../io/parquet/io_utils/parquet_io_utils.cpp | 92 ++++++++-- cpp/src/io/parquet/reader_impl_helpers.hpp | 1 + cpp/src/io/utilities/config_utils.cpp | 12 ++ cpp/tests/io/parquet_reader_test.cpp | 170 ++++++++++++++++++ 6 files changed, 290 insertions(+), 12 deletions(-) diff --git a/cpp/include/cudf/io/config_utils.hpp b/cpp/include/cudf/io/config_utils.hpp index d95afb1fa7dc..5f6612177304 100644 --- a/cpp/include/cudf/io/config_utils.hpp +++ b/cpp/include/cudf/io/config_utils.hpp @@ -1,11 +1,13 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include +#include + namespace CUDF_EXPORT cudf { namespace io { //! `KvikIO` @@ -67,5 +69,27 @@ namespace integrated_memory_optimization { /** @} */ // end of group } // namespace integrated_memory_optimization + +//! Parquet +namespace parquet_reader { + +/** + * @brief Returns the Parquet reader's footer speculative read size in bytes. + * + * 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(); + +} // namespace parquet_reader } // namespace io } // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/io/detail/parquet.hpp b/cpp/include/cudf/io/detail/parquet.hpp index 4e5c6c93a7e8..c9a43aa8ade9 100644 --- a/cpp/include/cudf/io/detail/parquet.hpp +++ b/cpp/include/cudf/io/detail/parquet.hpp @@ -18,6 +18,7 @@ #include +#include #include #include 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..e03f654ce890 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -10,6 +10,7 @@ #include #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,88 @@ 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_reader::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"); - CUDF_EXPECTS(ender->magic == detail::parquet_magic, "Corrupted footer"); - CUDF_EXPECTS(ender->footer_len != 0 && ender->footer_len <= (len - header_len - ender_len), + 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, + std::format("Failed to read Parquet speculative metadata bytes: " + "requested_offset={}, requested_size={}, bytes_read={}, " + "required_size={}", + speculative_read_offset, + speculative_read_size, + speculative_buffer->size(), + speculative_read_size)); + + file_ender_s ender{}; + std::memcpy( + &ender, speculative_buffer->data() + speculative_buffer->size() - ender_len, ender_len); + + if (speculative_read_offset == 0) { + file_header_s header{}; + std::memcpy(&header, speculative_buffer->data(), header_len); + 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) { + auto const footer_start_offset = footer_offset - speculative_read_offset; + CUDF_EXPECTS( + footer_start_offset + ender.footer_len <= speculative_buffer->size(), + std::format("Speculative metadata read did not include full footer bytes: " + "file_size={}, metadata_size_hint={}, speculative_read_offset={}, " + "speculative_read_size={}, bytes_read={}, footer_offset={}, footer_len={}", + len, + metadata_size_hint, + speculative_read_offset, + speculative_read_size, + speculative_buffer->size(), + footer_offset, + ender.footer_len)); + 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)); + } + + // Footer starts before the speculative read range. 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, + std::format("Failed to read the missing footer prefix bytes: " + "requested_offset={}, requested_size={}, bytes_read={}, file_size={}", + footer_offset, + missing_prefix_size, + missing_prefix->size(), + len)); + 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; + CUDF_EXPECTS(speculative_buffer->size() >= footer_suffix_size, + std::format("Failed to read Parquet speculative metadata suffix bytes: " + "requested_offset={}, requested_size={}, bytes_read={}, " + "required_size={}", + speculative_read_offset, + speculative_read_size, + speculative_buffer->size(), + footer_suffix_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) { @@ -338,8 +410,6 @@ fetch_byte_ranges_to_device_async_impl( 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/src/io/parquet/reader_impl_helpers.hpp b/cpp/src/io/parquet/reader_impl_helpers.hpp index a3238d66adbe..ed68c252baf5 100644 --- a/cpp/src/io/parquet/reader_impl_helpers.hpp +++ b/cpp/src/io/parquet/reader_impl_helpers.hpp @@ -15,6 +15,7 @@ #include #include +#include #include #include #include diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index f452babee592..a417c6c9acca 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -82,4 +82,16 @@ namespace integrated_memory_optimization { } } // namespace integrated_memory_optimization + +namespace parquet_reader { + +[[nodiscard]] std::size_t metadata_size_hint() +{ + static constexpr auto default_metadata_size_hint = std::size_t{64} * 1024; + static auto const metadata_size_hint = cudf::detail::getenv_or( + "LIBCUDF_PARQUET_METADATA_SIZE_HINT", default_metadata_size_hint); + return metadata_size_hint; +} + +} // namespace parquet_reader } // namespace cudf::io diff --git a/cpp/tests/io/parquet_reader_test.cpp b/cpp/tests/io/parquet_reader_test.cpp index 8ca9de9c10db..a92b72094778 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,109 @@ struct ParquetMetadataReaderTest : public cudf::test::BaseFixture { } }; +namespace { +class TrackingFooterDatasource : public cudf::io::datasource { + public: + explicit TrackingFooterDatasource(std::vector const& data) : data_(data) {} + + [[nodiscard]] std::vector> const& reads() const { return reads_; } + + void reset() { reads_.clear(); } + + std::unique_ptr host_read(size_t offset, size_t size) override + { + reads_.emplace_back(offset, size); + CUDF_EXPECTS(offset <= data_.size(), "Offset is out of bounds"); + auto const read_size = std::min(size, data_.size() - offset); + std::vector out(read_size); + std::memcpy(out.data(), data_.data() + offset, read_size); + return cudf::io::datasource::buffer::create(std::move(out)); + } + + size_t host_read(size_t offset, size_t size, uint8_t* dst) override + { + reads_.emplace_back(offset, size); + CUDF_EXPECTS(offset <= data_.size(), "Offset is out of bounds"); + auto const read_size = std::min(size, data_.size() - offset); + std::memcpy(dst, data_.data() + offset, read_size); + return read_size; + } + + [[nodiscard]] size_t size() const override { return data_.size(); } + + private: + std::vector const& data_; + std::vector> reads_; +}; + +class ShortReadFooterDatasource : public cudf::io::datasource { + public: + ShortReadFooterDatasource(std::vector const& data, + size_t short_read_call, + size_t short_read_size) + : data_(data), short_read_call_(short_read_call), short_read_size_(short_read_size) + { + } + + std::unique_ptr host_read(size_t offset, size_t size) override + { + ++call_count_; + CUDF_EXPECTS(offset <= data_.size(), "Offset is out of bounds"); + auto read_size = std::min(size, data_.size() - offset); + if (call_count_ == short_read_call_) { read_size = std::min(read_size, short_read_size_); } + std::vector out(read_size); + std::memcpy(out.data(), data_.data() + offset, read_size); + return cudf::io::datasource::buffer::create(std::move(out)); + } + + size_t host_read(size_t offset, size_t size, uint8_t* dst) override + { + ++call_count_; + CUDF_EXPECTS(offset <= data_.size(), "Offset is out of bounds"); + auto read_size = std::min(size, data_.size() - offset); + if (call_count_ == short_read_call_) { read_size = std::min(read_size, short_read_size_); } + std::memcpy(dst, data_.data() + offset, read_size); + return read_size; + } + + [[nodiscard]] size_t size() const override { return data_.size(); } + + private: + std::vector const& data_; + size_t short_read_call_; + size_t short_read_size_; + size_t call_count_{0}; +}; + +std::vector make_simple_parquet_bytes() +{ + auto ints = random_values(128); + cudf::test::fixed_width_column_wrapper int_col(ints.begin(), ints.end()); + cudf::table_view input_table({int_col}); + + std::vector parquet_bytes; + auto const write_opts = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&parquet_bytes}, input_table) + .build(); + cudf::io::write_parquet(write_opts); + return parquet_bytes; +} + +template +void expect_logic_error_contains(Fn&& fn, std::string const& needle) +{ + try { + fn(); + FAIL() << "Expected cudf::logic_error"; + } catch (cudf::logic_error const& e) { + EXPECT_NE(std::string{e.what()}.find(needle), std::string::npos) + << "Actual message: " << e.what(); + } catch (...) { + FAIL() << "Expected cudf::logic_error"; + } +} +} // namespace + TEST_F(ParquetMetadataReaderTest, Basics) { auto const num_rows = 1200; @@ -2969,6 +3075,70 @@ TEST_F(ParquetMetadataReaderTest, PreMaterializedMetadata) test_parquet_metadata(3); } +TEST_F(ParquetMetadataReaderTest, MetadataFooterErrorMessages) +{ + auto parquet_bytes = make_simple_parquet_bytes(); + + // Small source size check + { + std::vector tiny(8, '\0'); + TrackingFooterDatasource tiny_source(tiny); + auto const source = cudf::io::source_info{&tiny_source}; + expect_logic_error_contains([&] { (void)cudf::io::read_parquet_metadata(source); }, + "Incorrect data source"); + } + + // Speculative ender short read check + { + ShortReadFooterDatasource short_ender_source( + parquet_bytes, /*short_read_call=*/1, /*short_read_size=*/7); + auto const source = cudf::io::source_info{&short_ender_source}; + expect_logic_error_contains([&] { (void)cudf::io::read_parquet_metadata(source); }, + "Failed to read Parquet speculative metadata bytes"); + } + + // Header magic check when speculative read starts at offset 0 + { + auto bad_header = parquet_bytes; + bad_header[0] = 'B'; + bad_header[1] = 'A'; + bad_header[2] = 'D'; + bad_header[3] = '!'; + TrackingFooterDatasource bad_header_source(bad_header); + auto const source = cudf::io::source_info{&bad_header_source}; + expect_logic_error_contains([&] { (void)cudf::io::read_parquet_metadata(source); }, + "Corrupted header"); + } + + // Footer magic check + { + auto bad_footer_magic = parquet_bytes; + auto const footer_magic_offset = bad_footer_magic.size() - sizeof(uint32_t); + bad_footer_magic[footer_magic_offset + 0] = 'B'; + bad_footer_magic[footer_magic_offset + 1] = 'A'; + bad_footer_magic[footer_magic_offset + 2] = 'D'; + bad_footer_magic[footer_magic_offset + 3] = '!'; + TrackingFooterDatasource bad_footer_source(bad_footer_magic); + auto const source = cudf::io::source_info{&bad_footer_source}; + expect_logic_error_contains([&] { (void)cudf::io::read_parquet_metadata(source); }, + "Corrupted footer"); + } + + // Footer length check + { + auto bad_footer_len = parquet_bytes; + auto const footer_len_offset = bad_footer_len.size() - sizeof(cudf::io::parquet::file_ender_s); + bad_footer_len[footer_len_offset + 0] = 0; + bad_footer_len[footer_len_offset + 1] = 0; + bad_footer_len[footer_len_offset + 2] = 0; + bad_footer_len[footer_len_offset + 3] = 0; + TrackingFooterDatasource bad_footer_len_source(bad_footer_len); + auto const source = cudf::io::source_info{&bad_footer_len_source}; + expect_logic_error_contains([&] { (void)cudf::io::read_parquet_metadata(source); }, + "Incorrect footer length"); + } +} + TEST_F(ParquetMetadataReaderTest, Nested) { auto const num_rows = 1200; From b316e3d4fd124d335cf58c8057c9de3bbc584c26 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 9 Jun 2026 13:18:20 -0700 Subject: [PATCH 2/4] Merge tom/cudf-sourceinfo-size --- ci/build_wheel_cudf_streaming.sh | 6 +- cpp/include/cudf/io/datasource.hpp | 9 ++- cpp/include/cudf/io/types.hpp | 49 +++++++++++- cpp/src/io/functions.cpp | 18 +++-- cpp/src/io/utilities/datasource.cpp | 39 ++++++++-- cpp/tests/CMakeLists.txt | 1 + cpp/tests/io/filepath_source_test.cpp | 75 +++++++++++++++++++ docs/cudf/source/cudf/io/io.md | 23 ++++++ python/pylibcudf/pylibcudf/io/__init__.py | 3 +- python/pylibcudf/pylibcudf/io/types.pxd | 5 ++ python/pylibcudf/pylibcudf/io/types.pyi | 10 ++- python/pylibcudf/pylibcudf/io/types.pyx | 49 +++++++++++- .../pylibcudf/pylibcudf/libcudf/io/types.pxd | 14 +++- .../tests/io/test_source_sink_info.py | 42 ++++++++++- 14 files changed, 316 insertions(+), 27 deletions(-) create mode 100644 cpp/tests/io/filepath_source_test.cpp diff --git a/ci/build_wheel_cudf_streaming.sh b/ci/build_wheel_cudf_streaming.sh index c2619fe55798..c6c08919476b 100755 --- a/ci/build_wheel_cudf_streaming.sh +++ b/ci/build_wheel_cudf_streaming.sh @@ -12,11 +12,13 @@ dependency_file_key_suffix="cudf_streaming" RAPIDS_PY_CUDA_SUFFIX="$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")" -# Downloads libcudf_streaming wheel from this current build, -# then ensures 'cudf_streaming' wheel builds always use the 'libcudf_streaming' just built in the same CI run. +# Downloads libcudf, pylibcudf, and libcudf_streaming wheels from the current build. +# Then ensures 'cudf_streaming' wheel builds always use wheels built in the same CI run. LIBCUDF_STREAMING_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="libcudf_streaming_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github cpp) +LIBCUDF_WHEELHOUSE=$(RAPIDS_PY_WHEEL_NAME="libcudf_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-github cpp) PYLIBCUDF_WHEELHOUSE=$(rapids-download-from-github "$(rapids-package-name "wheel_python" pylibcudf --stable --cuda "$RAPIDS_CUDA_VERSION")") echo "libcudf-streaming-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${LIBCUDF_STREAMING_WHEELHOUSE}"/libcudf_streaming_*.whl)" >> "${PIP_CONSTRAINT}" +echo "libcudf-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${LIBCUDF_WHEELHOUSE}"/libcudf_*.whl)" >> "${PIP_CONSTRAINT}" echo "pylibcudf-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo "${PYLIBCUDF_WHEELHOUSE}"/pylibcudf_*.whl)" >> "${PIP_CONSTRAINT}" rapids-logger "Generating build requirements" diff --git a/cpp/include/cudf/io/datasource.hpp b/cpp/include/cudf/io/datasource.hpp index 6753c2638e39..a4a87bba180d 100644 --- a/cpp/include/cudf/io/datasource.hpp +++ b/cpp/include/cudf/io/datasource.hpp @@ -13,7 +13,7 @@ #include #include -#include +#include namespace CUDF_EXPORT cudf { //! IO interfaces @@ -98,11 +98,14 @@ class datasource { * @param[in] offset Starting byte offset from which data will be read (the default is zero) * @param[in] max_size_estimate Upper estimate of the data range that will be read (the default is * zero, which means the whole file after `offset`) + * @param[in] known_size Optional known file size in bytes. When set for remote URLs, the IO + * backend may skip querying the remote server for file size at open time. * @return Constructed datasource object */ static std::unique_ptr create(std::string const& filepath, - size_t offset = 0, - size_t max_size_estimate = 0); + size_t offset = 0, + size_t max_size_estimate = 0, + std::optional known_size = std::nullopt); /** * @brief Creates a source from a host memory buffer. diff --git a/cpp/include/cudf/io/types.hpp b/cpp/include/cudf/io/types.hpp index 4978960b7c30..d912740257ef 100644 --- a/cpp/include/cudf/io/types.hpp +++ b/cpp/include/cudf/io/types.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -310,6 +310,17 @@ constexpr inline auto is_byte_like_type() std::is_same_v; } +/** + * @brief A file path with an optional known size in bytes. + * + * When `size` is set for a remote URL, the IO backend may skip querying the remote server for file + * size at open time. + */ +struct filepath_source { + std::string path; ///< Path or URL of the input file + std::optional size{}; ///< Known file size; omit to query size at open time +}; + /** * @brief Source information for read interfaces */ @@ -325,8 +336,13 @@ struct source_info { * @param file_paths Input files paths */ explicit source_info(std::vector file_paths) - : _type(io_type::FILEPATH), _num_sources(file_paths.size()), _filepaths(std::move(file_paths)) + : _type(io_type::FILEPATH), _num_sources(file_paths.size()) { + _filepath_sources.reserve(file_paths.size()); + for (auto& path : file_paths) { + _filepath_sources.push_back({std::move(path), std::nullopt}); + } + rebuild_filepaths(); } /** @@ -335,10 +351,21 @@ struct source_info { * @param file_path Single input file */ explicit source_info(std::string file_path) - : _type(io_type::FILEPATH), _num_sources(1), _filepaths({std::move(file_path)}) + : source_info(std::vector{std::move(file_path)}) { } + /** + * @brief Construct a new source info object from filepath sources with optional known sizes + * + * @param sources Input filepath sources + */ + explicit source_info(std::vector sources) + : _type(io_type::FILEPATH), _num_sources(sources.size()), _filepath_sources(std::move(sources)) + { + rebuild_filepaths(); + } + /** * @brief Construct a new source info object for multiple buffers in host memory * @@ -424,6 +451,12 @@ struct source_info { * @return The type of the input */ [[nodiscard]] auto type() const { return _type; } + /** + * @brief Get the filepath sources of the input + * + * @return The filepath sources of the input + */ + [[nodiscard]] auto const& filepath_sources() const { return _filepath_sources; } /** * @brief Get the filepaths of the input * @@ -457,8 +490,18 @@ struct source_info { [[nodiscard]] auto num_sources() const { return _num_sources; } private: + void rebuild_filepaths() + { + _filepaths.clear(); + _filepaths.reserve(_filepath_sources.size()); + for (auto const& source : _filepath_sources) { + _filepaths.push_back(source.path); + } + } + io_type _type = io_type::VOID; size_t _num_sources = 0; + std::vector _filepath_sources; std::vector _filepaths; std::vector> _host_buffers; std::vector> _device_buffers; diff --git a/cpp/src/io/functions.cpp b/cpp/src/io/functions.cpp index 96dce2c1e2cb..eed9e71a90d1 100644 --- a/cpp/src/io/functions.cpp +++ b/cpp/src/io/functions.cpp @@ -162,24 +162,26 @@ std::vector> make_datasources(source_info switch (info.type()) { case io_type::FILEPATH: { std::vector> sources; - sources.reserve(info.filepaths().size()); + sources.reserve(info.filepath_sources().size()); // Creating sources in a single thread is faster for a small number of sources auto const pool_use_threshold = cudf::detail::getenv_or("LIBCUDF_DATASOURCE_PARALLEL_CREATION_THRESHOLD", 8ul); - if (info.filepaths().size() >= pool_use_threshold) { + if (info.filepath_sources().size() >= pool_use_threshold) { std::vector>> source_tasks; - source_tasks.reserve(info.filepaths().size()); - for (auto const& path : info.filepaths()) { - source_tasks.emplace_back(cudf::detail::host_worker_pool().submit_task( - [=] { return cudf::io::datasource::create(path, offset, max_size_estimate); })); + source_tasks.reserve(info.filepath_sources().size()); + for (auto const& fs : info.filepath_sources()) { + source_tasks.emplace_back(cudf::detail::host_worker_pool().submit_task([=] { + return cudf::io::datasource::create(fs.path, offset, max_size_estimate, fs.size); + })); } std::transform( source_tasks.begin(), source_tasks.end(), std::back_inserter(sources), [](auto& task) { return task.get(); }); } else { - for (auto const& filepath : info.filepaths()) { - sources.emplace_back(cudf::io::datasource::create(filepath, offset, max_size_estimate)); + for (auto const& fs : info.filepath_sources()) { + sources.emplace_back( + cudf::io::datasource::create(fs.path, offset, max_size_estimate, fs.size)); } } return sources; diff --git a/cpp/src/io/utilities/datasource.cpp b/cpp/src/io/utilities/datasource.cpp index aee5b7abfe52..fb6c49c81d8e 100644 --- a/cpp/src/io/utilities/datasource.cpp +++ b/cpp/src/io/utilities/datasource.cpp @@ -27,6 +27,7 @@ #include #ifdef CUDF_KVIKIO_REMOTE_IO +#include #include #endif @@ -354,13 +355,40 @@ class user_datasource_wrapper : public datasource { }; #ifdef CUDF_KVIKIO_REMOTE_IO +/** + * @brief Infer the KvikIO remote endpoint type from a URL (no network I/O). + * + * Mirrors the order used by `kvikio::RemoteHandle::open()` in AUTO mode. + */ +kvikio::RemoteEndpointType infer_remote_endpoint_type(std::string const& url) +{ + if (kvikio::S3Endpoint::is_url_valid(url)) { return kvikio::RemoteEndpointType::S3; } + if (kvikio::S3PublicEndpoint::is_url_valid(url)) { return kvikio::RemoteEndpointType::S3_PUBLIC; } + if (kvikio::S3EndpointWithPresignedUrl::is_url_valid(url)) { + return kvikio::RemoteEndpointType::S3_PRESIGNED_URL; + } + if (kvikio::WebHdfsEndpoint::is_url_valid(url)) { return kvikio::RemoteEndpointType::WEBHDFS; } + if (kvikio::HttpEndpoint::is_url_valid(url)) { return kvikio::RemoteEndpointType::HTTP; } + return kvikio::RemoteEndpointType::HTTP; +} + +kvikio::RemoteHandle open_remote_handle(char const* filepath, std::optional known_size) +{ + if (known_size.has_value()) { + auto const endpoint_type = infer_remote_endpoint_type(filepath); + return kvikio::RemoteHandle::open(filepath, endpoint_type, std::nullopt, *known_size); + } + return kvikio::RemoteHandle::open(filepath); +} + /** * @brief Remote file source backed by KvikIO, which handles S3 filepaths seamlessly. */ class remote_file_source : public kvikio_source { public: - explicit remote_file_source(char const* filepath) - : kvikio_source{kvikio::RemoteHandle::open(filepath)} + explicit remote_file_source(char const* filepath, + std::optional known_size = std::nullopt) + : kvikio_source{open_remote_handle(filepath, known_size)} { } @@ -397,7 +425,8 @@ class remote_file_source : public file_source { std::unique_ptr datasource::create(std::string const& filepath, size_t offset, - size_t max_size_estimate) + size_t max_size_estimate, + std::optional known_size) { auto const use_memory_mapping = [] { auto const policy = cudf::detail::getenv_or("LIBCUDF_MMAP_ENABLED", std::string{"OFF"}); @@ -410,7 +439,7 @@ std::unique_ptr datasource::create(std::string const& filepath, if (remote_file_source::could_be_remote_url(filepath)) { try { - return std::make_unique(filepath.c_str()); + return std::make_unique(filepath.c_str(), known_size); } catch (std::exception const& ex) { std::string redacted_msg; try { @@ -450,7 +479,7 @@ std::unique_ptr datasource::create(std::string const& filepath, // Create a remote file resource only when the pattern is found and replaced; otherwise, still // create a local file resource if (filepath != remote_file_path) { - return std::make_unique(remote_file_path.c_str()); + return std::make_unique(remote_file_path.c_str(), known_size); } } diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index e709a52ace20..cfc908341daa 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -320,6 +320,7 @@ ConfigureTest( # * io tests -------------------------------------------------------------------------------------- ConfigureTest(COMPRESSION_TEST io/comp/comp_test.cpp) ConfigureTest(ROW_SELECTION_TEST io/row_selection_test.cpp) +ConfigureTest(FILEPATH_SOURCE_TEST io/filepath_source_test.cpp) ConfigureTest( CSV_TEST io/csv_test.cpp diff --git a/cpp/tests/io/filepath_source_test.cpp b/cpp/tests/io/filepath_source_test.cpp new file mode 100644 index 000000000000..5eb11489615f --- /dev/null +++ b/cpp/tests/io/filepath_source_test.cpp @@ -0,0 +1,75 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include +#include +#include + +#include + +auto const temp_env = static_cast( + ::testing::AddGlobalTestEnvironment(new cudf::test::TempDirTestEnvironment)); + +struct FilepathSourceTest : public cudf::test::BaseFixture {}; + +TEST_F(FilepathSourceTest, StringConstructorsPopulateFilepathSources) +{ + auto const single = cudf::io::source_info{"test.parquet"}; + ASSERT_EQ(single.filepath_sources().size(), 1); + EXPECT_EQ(single.filepath_sources().front().path, "test.parquet"); + EXPECT_FALSE(single.filepath_sources().front().size.has_value()); + EXPECT_EQ(single.filepaths(), std::vector{"test.parquet"}); + + auto const multi = cudf::io::source_info{std::vector{"a.parquet", "b.parquet"}}; + ASSERT_EQ(multi.filepath_sources().size(), 2); + EXPECT_EQ(multi.filepaths().size(), 2); + EXPECT_FALSE(multi.filepath_sources()[1].size.has_value()); +} + +TEST_F(FilepathSourceTest, FilepathSourceConstructorPreservesSize) +{ + std::vector sources{ + {"s3://bucket/object.parquet", 12345}, + {"https://example.com/data.parquet", std::nullopt}, + }; + + auto const info = cudf::io::source_info{std::move(sources)}; + ASSERT_EQ(info.filepath_sources().size(), 2); + EXPECT_EQ(info.filepath_sources()[0].path, "s3://bucket/object.parquet"); + ASSERT_TRUE(info.filepath_sources()[0].size.has_value()); + EXPECT_EQ(info.filepath_sources()[0].size.value(), 12345); + EXPECT_FALSE(info.filepath_sources()[1].size.has_value()); + EXPECT_EQ(info.filepaths()[0], "s3://bucket/object.parquet"); + EXPECT_EQ(info.filepaths()[1], "https://example.com/data.parquet"); +} + +TEST_F(FilepathSourceTest, KnownSizePlumbsThroughMakeDatasources) +{ + auto const filepath = temp_env->get_temp_filepath("KnownSize.parquet"); + + auto col = cudf::test::fixed_width_column_wrapper{1, 2, 3}; + cudf::table_view const table{{col}}; + + cudf::io::parquet_writer_options write_opts = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, table); + cudf::io::write_parquet(write_opts); + + auto const file_size = std::filesystem::file_size(filepath); + std::vector sources{{filepath, file_size}}; + auto const source_info = cudf::io::source_info{std::move(sources)}; + + auto datasources = cudf::io::make_datasources(source_info); + ASSERT_EQ(datasources.size(), 1); + EXPECT_EQ(datasources.front()->size(), file_size); + + 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(table, result.tbl->view()); +} diff --git a/docs/cudf/source/cudf/io/io.md b/docs/cudf/source/cudf/io/io.md index e46f5d09a973..688ef87fde0e 100644 --- a/docs/cudf/source/cudf/io/io.md +++ b/docs/cudf/source/cudf/io/io.md @@ -117,6 +117,29 @@ Note that: For more information about error handling, compatibility mode, and tuning parameters in KvikIO see: +### Remote file sizes and HEAD requests + +When reading remote files (for example `s3://...` URLs) via `pylibcudf.io.SourceInfo`, KvikIO +may send HEAD requests at open time to probe connectivity and query file size. To skip those +requests when the file size is already known (for example from object-store metadata), pass a +`pylibcudf.io.FilepathSource` with the `size` argument set: + +```python +import pylibcudf as plc + +content_length = ... # from external metadata +sources = plc.io.SourceInfo([ + plc.io.FilepathSource("s3://bucket/object.parquet", size=content_length), +]) +table = plc.io.parquet.read_parquet( + plc.io.parquet.ParquetReaderOptions.builder(sources).build() +) +``` + +Providing an incorrect size avoids the extra HEAD requests but will break footer reads and other +operations that depend on the true file length. Plain string paths in `SourceInfo` preserve the +previous behavior (size queried via KvikIO). + Operations that support the use of GPUDirect Storage: - {py:func}`cudf.read_avro` diff --git a/python/pylibcudf/pylibcudf/io/__init__.py b/python/pylibcudf/pylibcudf/io/__init__.py index 2162b50e963c..a256857920c9 100644 --- a/python/pylibcudf/pylibcudf/io/__init__.py +++ b/python/pylibcudf/pylibcudf/io/__init__.py @@ -15,10 +15,11 @@ types, ) from .parquet_metadata import FileMetaData -from .types import SinkInfo, SourceInfo, TableWithMetadata +from .types import FilepathSource, SinkInfo, SourceInfo, TableWithMetadata __all__ = [ "FileMetaData", + "FilepathSource", "SinkInfo", "SourceInfo", "TableWithMetadata", diff --git a/python/pylibcudf/pylibcudf/io/types.pxd b/python/pylibcudf/pylibcudf/io/types.pxd index 1e52f4faa058..5b476d066627 100644 --- a/python/pylibcudf/pylibcudf/io/types.pxd +++ b/python/pylibcudf/pylibcudf/io/types.pxd @@ -14,6 +14,7 @@ from pylibcudf.libcudf.io.types cimport ( column_name_info, compression_type, dictionary_policy, + filepath_source, io_type, partition_info, quote_style, @@ -88,6 +89,10 @@ cdef class TableWithMetadata: table_with_metadata& tbl, object stream, DeviceMemoryResource mr ) +cdef class FilepathSource: + cdef public object path + cdef public object size + cdef class SourceInfo: cdef source_info c_obj # Keep the bytes converted from stringio alive diff --git a/python/pylibcudf/pylibcudf/io/types.pyi b/python/pylibcudf/pylibcudf/io/types.pyi index f2050a5b1f91..0d491524b5c0 100644 --- a/python/pylibcudf/pylibcudf/io/types.pyi +++ b/python/pylibcudf/pylibcudf/io/types.pyi @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 import io import os @@ -111,11 +111,19 @@ class TableWithMetadata: @property def num_row_groups_after_bloom_filter(self) -> int | None: ... +class FilepathSource: + def __init__( + self, path: str | os.PathLike[Any], size: int | None = None + ): ... + path: str + size: int | None + class SourceInfo: def __init__( self, sources: Sequence[str] | Sequence[os.PathLike[Any]] + | Sequence[FilepathSource] | Sequence[Datasource], ) -> None: ... @staticmethod diff --git a/python/pylibcudf/pylibcudf/io/types.pyx b/python/pylibcudf/pylibcudf/io/types.pyx index 27c3bb47caf0..6805f635f9ce 100644 --- a/python/pylibcudf/pylibcudf/io/types.pyx +++ b/python/pylibcudf/pylibcudf/io/types.pyx @@ -22,6 +22,7 @@ from pylibcudf.libcudf.io.types cimport ( column_encoding, column_in_metadata, column_name_info, + filepath_source, partition_info, source_info, table_input_metadata, @@ -55,6 +56,7 @@ __all__ = [ "ColumnInMetadata", "CompressionType", "DictionaryPolicy", + "FilepathSource", "JSONRecoveryMode", "PartitionInfo", "QuoteStyle", @@ -453,6 +455,27 @@ cdef class TableWithMetadata: return None +cdef class FilepathSource: + """ + A file path or URL with an optional known size in bytes. + + When ``size`` is set for a remote URL, libcudf passes it to KvikIO at open + time so the remote server is not queried for file size (avoiding HEAD + requests). An incorrect size will cause read failures. + + Parameters + ---------- + path : str or os.PathLike + Path or URL of the input file. + size : int, optional + Known file size in bytes. Omit to query size via KvikIO (HEAD for remote URLs). + """ + + def __init__(self, path, size=None): + self.path = os.fspath(path) + self.size = size + + cdef class SourceInfo: """ A class containing details on a source to read from. @@ -464,6 +487,7 @@ cdef class SourceInfo: sources : List[Union[ str, os.PathLike, + FilepathSource, bytes, io.BytesIO, DataSource, @@ -480,9 +504,30 @@ cdef class SourceInfo: return cdef vector[string] c_files + cdef vector[filepath_source] c_filepath_sources cdef vector[datasource*] c_datasources + cdef filepath_source fs + + if isinstance(sources[0], FilepathSource): + c_filepath_sources.reserve(len(sources)) + + for src in sources: + if not isinstance(src, FilepathSource): + raise ValueError("All sources must be of the same type!") + if not ( + os.path.isfile(src.path) or SourceInfo._is_remote_uri(src.path) + ): + raise FileNotFoundError( + errno.ENOENT, os.strerror(errno.ENOENT), src.path + ) + fs = filepath_source( str(src.path).encode()) + if src.size is not None: + fs.size = src.size + c_filepath_sources.push_back(fs) - if isinstance(sources[0], (os.PathLike, str)): + self.c_obj = move(source_info(c_filepath_sources)) + return + elif isinstance(sources[0], (os.PathLike, str)): c_files.reserve(len(sources)) for src in sources: @@ -537,7 +582,7 @@ cdef class SourceInfo: self.c_obj = move(source_info(host_span[device_span[const_byte]](d_spans))) return else: - raise ValueError("Sources must be a list of str/paths, " + raise ValueError("Sources must be a list of str/paths, FilepathSource, " "bytes, io.BytesIO, io.StringIO, or a Datasource") self.c_obj = source_info(host_span[host_span[const_byte]](self._hspans)) diff --git a/python/pylibcudf/pylibcudf/libcudf/io/types.pxd b/python/pylibcudf/pylibcudf/libcudf/io/types.pxd index 6a6d6356801e..938154f8be0f 100644 --- a/python/pylibcudf/pylibcudf/libcudf/io/types.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/io/types.pxd @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 cimport pylibcudf.libcudf.io.data_sink as cudf_io_data_sink cimport pylibcudf.libcudf.io.datasource as cudf_io_datasource @@ -125,13 +125,25 @@ cdef extern from "cudf/io/types.hpp" \ size_type start_row, size_type num_rows ) except +libcudf_exception_handler + cdef cppclass filepath_source: + string path + optional[size_t] size + + filepath_source() except +libcudf_exception_handler + filepath_source(string path) except +libcudf_exception_handler + cdef cppclass source_info: const vector[string]& filepaths() except +libcudf_exception_handler + const vector[filepath_source]& filepath_sources() \ + except +libcudf_exception_handler source_info() except +libcudf_exception_handler source_info( const vector[string] &filepaths ) except +libcudf_exception_handler + source_info( + vector[filepath_source] sources + ) except +libcudf_exception_handler source_info( cudf_io_datasource.datasource *source ) except +libcudf_exception_handler diff --git a/python/pylibcudf/tests/io/test_source_sink_info.py b/python/pylibcudf/tests/io/test_source_sink_info.py index 5a2bc95bd109..0bf2091be514 100644 --- a/python/pylibcudf/tests/io/test_source_sink_info.py +++ b/python/pylibcudf/tests/io/test_source_sink_info.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 import io @@ -93,3 +93,43 @@ def test_source_info_ctor_mixing_invalid(io_class, sources, tmp_path): def test_source_info_invalid(): with pytest.raises(ValueError): plc.io.SourceInfo([123]) + + +def test_filepath_source_local_parquet(tmp_path): + path = tmp_path / "data.parquet" + table = plc.Table([plc.Column.from_iterable_of_py([1, 2, 3])]) + plc.io.parquet.write_parquet( + plc.io.parquet.ParquetWriterOptions.builder( + plc.io.SinkInfo([str(path)]), table + ).build() + ) + file_size = path.stat().st_size + + source = plc.io.FilepathSource(str(path), size=file_size) + assert source.path == str(path) + assert source.size == file_size + + source_info = plc.io.SourceInfo([source]) + read_opts = plc.io.parquet.ParquetReaderOptions.builder( + source_info + ).build() + result = plc.io.parquet.read_parquet(read_opts) + assert result.columns[0].to_arrow().to_pylist() == [1, 2, 3] + + +def test_filepath_source_remote_uri_without_size(): + source = plc.io.FilepathSource("s3://bucket/object.parquet") + assert source.size is None + plc.io.SourceInfo([source]) + + +def test_filepath_source_mixed_sources_invalid(): + with pytest.raises( + ValueError, match="All sources must be of the same type" + ): + plc.io.SourceInfo( + [ + plc.io.FilepathSource("s3://bucket/object.parquet", size=100), + "s3://bucket/other.parquet", + ] + ) From ece723ec659232838436e9210af26e2e9bade6a3 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 9 Jun 2026 13:19:49 -0700 Subject: [PATCH 3/4] Merge tom/cudf-polars-prefetch-metadata --- python/cudf_polars/cudf_polars/callback.py | 10 +- python/cudf_polars/cudf_polars/dsl/ir.py | 158 ++++++++++++++++-- python/cudf_polars/cudf_polars/engine/core.py | 27 +-- .../cudf_polars/cudf_polars/streaming/io.py | 41 +++-- .../cudf_polars/streaming/select.py | 9 +- .../cudf_polars/cudf_polars/utils/config.py | 13 ++ .../cudf_polars/tests/streaming/test_scan.py | 127 +++++++++++++- python/cudf_polars/tests/test_config.py | 3 + python/cudf_polars/tests/test_scan.py | 49 ++++++ python/cudf_polars/tests/test_select.py | 83 +++++++++ 10 files changed, 479 insertions(+), 41 deletions(-) diff --git a/python/cudf_polars/cudf_polars/callback.py b/python/cudf_polars/cudf_polars/callback.py index 93dd0790bd59..533eb8245623 100644 --- a/python/cudf_polars/cudf_polars/callback.py +++ b/python/cudf_polars/cudf_polars/callback.py @@ -24,7 +24,10 @@ from rmm._cuda import gpu import cudf_polars.dsl.tracing -from cudf_polars.dsl.ir import IRExecutionContext +from cudf_polars.dsl.ir import ( + IRExecutionContext, + prefetch_parquet_file_metadata_for_ir, +) from cudf_polars.dsl.tracing import CUDF_POLARS_NVTX_DOMAIN from cudf_polars.dsl.translate import Translator from cudf_polars.utils.config import ( @@ -301,6 +304,11 @@ def _callback( ): if config_options.executor.name == "in-memory": context = IRExecutionContext() + if config_options.parquet_options.prefetch_file_metadata: + prefetch_parquet_file_metadata_for_ir( + ir, + context, + ) df = ir.evaluate(cache={}, timer=timer, context=context).to_polars() if timer is None: return df diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index cb55d9c01308..06058867eaa7 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -14,6 +14,7 @@ from __future__ import annotations import asyncio +import concurrent.futures import contextlib import contextvars import functools @@ -47,7 +48,11 @@ from cudf_polars.dsl.expressions.base import ExecutionContext from cudf_polars.dsl.nodebase import Node from cudf_polars.dsl.to_ast import _DECIMAL_IDS, to_ast, to_parquet_filter -from cudf_polars.dsl.tracing import log_do_evaluate, nvtx_annotate_cudf_polars +from cudf_polars.dsl.tracing import ( + log_do_evaluate, + nvtx_annotate_cudf_polars, +) +from cudf_polars.dsl.traversal import traversal from cudf_polars.dsl.utils.reshape import broadcast from cudf_polars.dsl.utils.windows import ( offsets_to_windows, @@ -66,7 +71,6 @@ if TYPE_CHECKING: from collections.abc import Callable, Generator, Hashable, Iterable, Sequence - from concurrent.futures import ThreadPoolExecutor from typing import Literal, Self from polars import polars # type: ignore[attr-defined] @@ -127,11 +131,18 @@ class IRExecutionContext: A zero-argument callable that returns a CUDA stream. query_id Identifier for the query being executed. + parquet_file_metadata + A cache of parquet file metadata. The keys are the ``paths`` of Scan nodes + with a ``parquet`` type. The values are a list of ``FileMetaData`` objects + associated with those ``paths``. """ - py_executor: ThreadPoolExecutor | None = field(default=None) + py_executor: concurrent.futures.ThreadPoolExecutor | None = field(default=None) get_cuda_stream: Callable[[], Stream] = field(default=get_cuda_stream) query_id: uuid.UUID = field(default_factory=uuid.uuid4) + parquet_file_metadata: dict[ + tuple[str, ...], list[plc.io.parquet_metadata.FileMetaData] + ] = field(default_factory=dict) async def to_thread( self, func: Callable[P, T], /, *args: P.args, **kwargs: P.kwargs @@ -185,6 +196,92 @@ def stream_ordered_after(self, *dfs: DataFrame) -> Generator[Stream, None, None] yield result_stream +@nvtx_annotate_cudf_polars(message="fetch_parquet_footers_for_paths") +def _prefetch_parquet_footers_for_paths( + paths: tuple[str, ...], +) -> tuple[tuple[str, ...], list[plc.io.parquet_metadata.FileMetaData]]: + """ + Prefetch parquet footers for a list of paths. + + This is typically executed concurrently with prefetch operations for other + groups of ``paths`` for other ``Scan`` nodes. + + Parameters + ---------- + paths + The tuple of paths to prefetch. These correspond to ``paths`` in a ``Scan`` node. + + Returns + ------- + paths + The original input ``paths``. Useful for associating the result with the metadata + when executing out of order concurrently. + metadata + The list of ``FileMetaData`` objects for the ``paths``. + """ + metadata = plc.io.parquet_metadata.read_parquet_footers( + plc.io.SourceInfo(list(paths)) + ) + return paths, metadata + + +@nvtx_annotate_cudf_polars(message="prefetch_parquet_file_metadata_for_ir") +def prefetch_parquet_file_metadata_for_ir( + root: IR, + context: IRExecutionContext, +) -> None: + """ + Prefetch parquet metadata for all parquet scans in an IR graph. + + Parameters + ---------- + root + The root of the IR graph, which will be traversed. + context + The IR execution context. Its ``py_executor`` is used to fetch + metadata concurrently, its ``parquet_file_metadata`` is mutated + to cache the newly read parquet metadata. + """ + from cudf_polars.streaming.io import SplitScan, StreamingScan + + groups = set() + for node in traversal([root]): + if isinstance(node, StreamingScan): + for scan in node.scans: + if isinstance(scan, Scan) and scan.typ == "parquet": + groups.add(tuple(scan.paths)) + elif isinstance(scan, SplitScan) and scan.base_scan.typ == "parquet": + groups.add(tuple(scan.base_scan.paths)) + elif isinstance(node, Scan) and node.typ == "parquet": + groups.add(tuple(node.paths)) + + if not groups: + return + + missing_paths = { + paths for paths in groups if paths not in context.parquet_file_metadata + } + cm: contextlib.AbstractContextManager[concurrent.futures.Executor | None] + + if context.py_executor is None: + cm = executor = concurrent.futures.ThreadPoolExecutor() + else: + executor = context.py_executor + # We didn't create the executor, so we don't close it. + cm = contextlib.nullcontext() + + if missing_paths: + with cm: + futures = [ + executor.submit(_prefetch_parquet_footers_for_paths, paths) + for paths in missing_paths + ] + + for future in concurrent.futures.as_completed(futures): + paths, metadata = future.result() + context.parquet_file_metadata.setdefault(paths, metadata) + + _BINOPS = { plc.binaryop.BinaryOperator.EQUAL, plc.binaryop.BinaryOperator.NOT_EQUAL, @@ -352,6 +449,7 @@ def __init__(self, schema: Schema, options: Any, predicate: expr.NamedExpr | Non def _parquet_physical_types( paths: list[str], columns: list[str] | None ) -> dict[str, plc.DataType]: + # This may not be able use prefetched metadata, since we don't (currently) have a Schema. metadata = plc.io.parquet_metadata.read_parquet_metadata(plc.io.SourceInfo(paths)) column_types = metadata.schema().column_types() @@ -648,12 +746,32 @@ def add_file_paths( @staticmethod @nvtx_annotate_cudf_polars(message="Scan._get_parquet_row_count_from_metadata") def _get_parquet_row_count_from_metadata( - paths: list[str], skip_rows: int, n_rows: int + paths: list[str], + skip_rows: int, + n_rows: int, + parquet_options: ParquetOptions, + context: IRExecutionContext | None, ) -> int: # Zero-width parquet files lose their row count when read through # pylibcudf. See https://github.com/rapidsai/cudf/issues/21428 - meta = plc.io.parquet_metadata.read_parquet_metadata(plc.io.SourceInfo(paths)) - num_rows = meta.num_rows() - skip_rows + if parquet_options.prefetch_file_metadata and context is not None: + try: + parquet_metadatas = context.parquet_file_metadata[tuple(paths)] + except KeyError as e: + msg = ( + f"Parquet file metadata was not prefetched for paths: {list(paths)}." + "Please report this as a bug to cudf-polars. You can work around it " + "by setting 'CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA=0'." + ) + raise AssertionError(msg) from e + num_rows = sum(metadata.num_rows for metadata in parquet_metadatas) + else: + meta = plc.io.parquet_metadata.read_parquet_metadata( + plc.io.SourceInfo(paths) + ) + num_rows = meta.num_rows() + + num_rows -= skip_rows if n_rows != -1: num_rows = min(num_rows, n_rows) return max(num_rows, 0) @@ -789,6 +907,19 @@ def read_csv_header( df, ) elif typ == "parquet": + if parquet_options.prefetch_file_metadata: + try: + parquet_metadatas = context.parquet_file_metadata[tuple(paths)] + except KeyError as e: + msg = ( + f"Parquet file metadata was not prefetched for paths: {list(paths)}." + "Please report this as a bug to cudf-polars. You can work around it " + "by setting 'CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA=0'." + ) + raise AssertionError(msg) from e + else: + parquet_metadatas = None + filters = None if predicate is not None and row_index is None: # Can't apply filters during read if we have a row index. @@ -817,6 +948,7 @@ def read_csv_header( parquet_reader_options, chunk_read_limit=parquet_options.chunk_read_limit, pass_read_limit=parquet_options.pass_read_limit, + parquet_metadatas=parquet_metadatas, stream=stream, ) chunk = reader.read_chunk() @@ -832,7 +964,9 @@ def read_csv_header( [concatenated_columns[i], columns.pop()], stream=stream ) num_rows = ( - cls._get_parquet_row_count_from_metadata(paths, skip_rows, n_rows) + cls._get_parquet_row_count_from_metadata( + paths, skip_rows, n_rows, parquet_options, context + ) if not names else None ) @@ -849,12 +983,16 @@ def read_csv_header( ) else: tbl_w_meta = plc.io.parquet.read_parquet( - parquet_reader_options, stream=stream + parquet_reader_options, + parquet_metadatas=parquet_metadatas, + stream=stream, ) # TODO: consider nested column names? col_names = tbl_w_meta.column_names(include_children=False) num_rows = ( - cls._get_parquet_row_count_from_metadata(paths, skip_rows, n_rows) + cls._get_parquet_row_count_from_metadata( + paths, skip_rows, n_rows, parquet_options, context + ) if not col_names else None ) @@ -1532,7 +1670,7 @@ def evaluate( stream = context.get_cuda_stream() scan = self.children[0] effective_rows = Scan._get_parquet_row_count_from_metadata( - scan.paths, scan.skip_rows, scan.n_rows + scan.paths, scan.skip_rows, scan.n_rows, scan.parquet_options, context ) dtype = DataType(pl.UInt32()) col = Column( diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index bfe97050ccf3..5d89b09cf097 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -25,7 +25,7 @@ from rapidsmpf.streaming.core.actor import run_actor_network from cudf_polars.containers import DataFrame -from cudf_polars.dsl.ir import IRExecutionContext +from cudf_polars.dsl.ir import IRExecutionContext, prefetch_parquet_file_metadata_for_ir from cudf_polars.streaming.actor_graph.collectives import ReserveOpIDs from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id from cudf_polars.streaming.actor_graph.core import generate_network @@ -408,14 +408,12 @@ def _find_memory_error(exc: BaseException) -> MemoryError | None: def execute_ir_on_rank( ctx: Context, comm: Communicator, - py_executor: ThreadPoolExecutor, ir: IR, + ir_context: IRExecutionContext, partition_info: MutableMapping[IR, PartitionInfo], config_options: ConfigOptions[StreamingExecutor], stats: StatsCollector, collective_id_map: dict[IR, list[int]], - *, - query_id: uuid.UUID, ) -> tuple[pl.DataFrame, list[ChannelMetadata]]: """ Execute a Polars IR query on a single rank's GPU. @@ -430,10 +428,10 @@ def execute_ir_on_rank( The active RapidsMPF streaming context for this rank. comm The active RapidsMPF communicator for this rank. - py_executor - Thread-pool executor used to drive the actor network. ir Root IR node describing the query. + ir_context + Execution context reused across scan-task execution. partition_info Per-node partition metadata. config_options @@ -442,8 +440,6 @@ def execute_ir_on_rank( Statistics collector. collective_id_map Mapping from IR nodes to their pre-allocated collective operation IDs. - query_id - Unique identifier for the query, propagated into actor traces. Returns ------- @@ -452,9 +448,6 @@ def execute_ir_on_rank( metadata Collected channel metadata. """ - ir_context = IRExecutionContext( - py_executor, get_cuda_stream=ctx.get_stream_from_pool, query_id=query_id - ) metadata_collector: list[ChannelMetadata] = [] nodes, output = generate_network( @@ -695,15 +688,23 @@ def evaluate_on_rank( # so we only log it once. log_query_plan(ir, config_options) + ir_context = IRExecutionContext( + py_executor, get_cuda_stream=ctx.get_stream_from_pool, query_id=query_id + ) + if config_options.parquet_options.prefetch_file_metadata: + prefetch_parquet_file_metadata_for_ir( + ir, + ir_context, + ) + with ReserveOpIDs(ir, config_options) as collective_id_map: return execute_ir_on_rank( ctx, comm, - py_executor, ir, + ir_context, partition_info, config_options, stats, collective_id_map, - query_id=query_id, ) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index d56c31ce379b..34bcf34bac01 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -306,10 +306,32 @@ def do_evaluate( # - We can use all this information to calculate the # "skip_rows" and "n_rows" options to use locally. - rowgroup_metadata = plc.io.parquet_metadata.read_parquet_metadata( - plc.io.SourceInfo(paths) - ).rowgroup_metadata() - total_row_groups = len(rowgroup_metadata) + if parquet_options.prefetch_file_metadata: + try: + parquet_metadatas = context.parquet_file_metadata[tuple(paths)] + except KeyError as e: + msg = ( + f"Parquet file metadata was not prefetched for paths: {list(paths)}." + "Please report this as a bug to cudf-polars. You can work around it " + "by setting 'CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA=0'." + ) + raise AssertionError(msg) from e + + row_group_num_rows = [ + num_rows + for metadata in parquet_metadatas + for num_rows in metadata.row_group_num_rows + ] + + else: + row_group_num_rows = [ + rg["num_rows"] + for rg in plc.io.parquet_metadata.read_parquet_metadata( + plc.io.SourceInfo(paths) + ).rowgroup_metadata() + ] + + total_row_groups = len(row_group_num_rows) if total_splits <= total_row_groups: # We have enough row-groups in the file to align # all "total_splits" of our reads with row-group @@ -318,17 +340,14 @@ def do_evaluate( # the row-group indices to "skip_rows" and "n_rows". rg_stride = total_row_groups // total_splits skip_rgs = rg_stride * split_index - skip_rows = sum(rg["num_rows"] for rg in rowgroup_metadata[:skip_rgs]) - n_rows = sum( - rg["num_rows"] - for rg in rowgroup_metadata[skip_rgs : skip_rgs + rg_stride] - ) + skip_rows = sum(row_group_num_rows[:skip_rgs]) + n_rows = sum(row_group_num_rows[skip_rgs : skip_rgs + rg_stride]) else: # There are not enough row-groups to align # all "total_splits" of our reads with row-group # boundaries. Use metadata to directly calculate # "skip_rows" and "n_rows" for the current read. - total_rows = sum(rg["num_rows"] for rg in rowgroup_metadata) + total_rows = sum(row_group_num_rows) n_rows = total_rows // total_splits skip_rows = n_rows * split_index @@ -424,7 +443,7 @@ def can_use_native_parquet_node( @lower_ir_node.register(Scan) def _( ir: Scan, rec: LowerIRTransformer -) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: +) -> tuple[StreamingScan, MutableMapping[IR, PartitionInfo]]: config_options = rec.state["config_options"] parquet_options = config_options.parquet_options if ( diff --git a/python/cudf_polars/cudf_polars/streaming/select.py b/python/cudf_polars/cudf_polars/streaming/select.py index e1466f0b6610..81bfec030bc9 100644 --- a/python/cudf_polars/cudf_polars/streaming/select.py +++ b/python/cudf_polars/cudf_polars/streaming/select.py @@ -431,8 +431,15 @@ def _( if scan_child and scan_child.predicate is None and scan_child.typ == "parquet": # Special Case: Fast count. + # We can't use prefetched file metadata here, because we're in lowering, + # not execution, so we don't have an IRExecutionContext with the prefetched + # file metadata yet. count = Scan._get_parquet_row_count_from_metadata( - scan_child.paths, scan_child.skip_rows, scan_child.n_rows + scan_child.paths, + scan_child.skip_rows, + scan_child.n_rows, + scan_child.parquet_options, + context=None, ) dtype = ir.exprs[0].value.dtype diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 5786a5351cc4..8c7282991967 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -206,6 +206,10 @@ class ParquetOptions: Whether to use the native rapidsmpf node for parquet reading. This option is only used by the streaming executor. Default is False. + prefetch_file_metadata + Whether to prefetch parquet file metadata and pass it through + `parquet_metadatas` to avoid rereading file footers. + Default is False. """ _env_prefix = "CUDF_POLARS__PARQUET_OPTIONS" @@ -247,6 +251,13 @@ class ParquetOptions: default=False, ) ) + prefetch_file_metadata: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__PREFETCH_FILE_METADATA", + _bool_converter, + default=False, + ) + ) def __post_init__(self) -> None: # noqa: D105 if not isinstance(self.chunked, bool): @@ -263,6 +274,8 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("max_row_group_samples must be an int") if not isinstance(self.use_rapidsmpf_native, bool): raise TypeError("use_rapidsmpf_native must be a bool") + if not isinstance(self.prefetch_file_metadata, bool): + raise TypeError("prefetch_file_metadata must be a bool") def default_target_partition_size(min_device_size: int | None) -> int: diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index c9ccb13202dc..675003e634cf 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -11,7 +11,12 @@ from cudf_polars import Translator from cudf_polars.containers import DataType -from cudf_polars.dsl.ir import IRExecutionContext, Scan +from cudf_polars.dsl.ir import ( + Empty, + IRExecutionContext, + Scan, + prefetch_parquet_file_metadata_for_ir, +) from cudf_polars.engine.options import StreamingOptions from cudf_polars.streaming.base import IOPartitionFlavor, IOPartitionPlan from cudf_polars.streaming.io import SplitScan, StreamingScan, expand_scan_for_rank @@ -63,6 +68,56 @@ def test_scan_parquet_use_rapidsmpf_native(tmp_path, df, streaming_engine_factor assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) +def test_scan_parquet_prefetch_file_metadata(tmp_path, df, streaming_engine_factory): + streaming_engine = streaming_engine_factory( + StreamingOptions( + target_partition_size=1_000, + parquet_options={"prefetch_file_metadata": True}, + ), + ) + make_partitioned_source(df, tmp_path, "parquet", n_files=2) + assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) + + +def test_scan_parquet_prefetch_file_metadata_fused_files( + tmp_path, df, streaming_engine_factory +) -> None: + streaming_engine = streaming_engine_factory( + StreamingOptions( + target_partition_size=1_000_000, + parquet_options={"prefetch_file_metadata": True}, + ), + ) + make_partitioned_source(df, tmp_path, "parquet", n_files=3) + assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) + + +def test_scan_parquet_prefetch_file_metadata_split_files( + tmp_path, df, streaming_engine_factory +) -> None: + streaming_engine = streaming_engine_factory( + StreamingOptions( + target_partition_size=1_000, + parquet_options={"prefetch_file_metadata": True}, + ), + ) + make_partitioned_source(df, tmp_path, "parquet", n_files=1) + assert_gpu_result_equal(pl.scan_parquet(tmp_path), engine=streaming_engine) + + +def test_prefetch_file_metadata_non_parquet_scan(df, streaming_engine_factory) -> None: + streaming_engine = streaming_engine_factory( + StreamingOptions(parquet_options={"prefetch_file_metadata": True}), + ) + assert_gpu_result_equal(df.lazy().select("x"), engine=streaming_engine) + + +def test_prefetch_parquet_file_metadata_no_parquet_scans() -> None: + context = IRExecutionContext() + prefetch_parquet_file_metadata_for_ir(Empty({}), context) + assert context.parquet_file_metadata == {} + + # --------------------------------------------------------------------------- # Tests migrated from tests/streaming/test_scan.py # --------------------------------------------------------------------------- @@ -158,7 +213,10 @@ def test_scan_union(engine: pl.GPUEngine, tmp_path: Path) -> None: assert_gpu_result_equal(q, engine=engine) -def _make_parquet_scan(paths: list[str]) -> Scan: +def _make_parquet_scan( + paths: list[str], parquet_options: ParquetOptions | None = None +) -> Scan: + parquet_options = parquet_options or ParquetOptions() return Scan( {"x": DataType(pl.Int64())}, "parquet", @@ -171,7 +229,7 @@ def _make_parquet_scan(paths: list[str]) -> Scan: None, None, None, - ParquetOptions(), + parquet_options, ) @@ -248,9 +306,68 @@ def test_expand_scan_for_rank_split_files( assert scan.base_scan.paths == ["file.parquet"] -def test_streaming_scan_raises() -> None: +def test_scan_missing_prefetch_metadata_raises() -> None: + # This isn't reachable by normal cudf-polars usage. + scan = _make_parquet_scan( + ["file.parquet"], parquet_options=ParquetOptions(prefetch_file_metadata=True) + ) + ctx = IRExecutionContext() + with pytest.raises( + AssertionError, + match=r"Parquet file metadata was not prefetched for paths: \['file\.parquet'\]\.", + ): + Scan.do_evaluate( + scan.schema, + scan.typ, + scan.reader_options, + scan.paths, + scan.with_columns, + scan.skip_rows, + scan.n_rows, + scan.row_index, + scan.include_file_paths, + scan.predicate, + scan.parquet_options, + context=ctx, + ) + + +def test_streaming_scan_missing_prefetch_metadata_raises() -> None: # This isn't reachable by normal cudf-polars usage. - scan = _make_parquet_scan(["file.parquet"]) + scan = _make_parquet_scan( + ["file.parquet"], parquet_options=ParquetOptions(prefetch_file_metadata=True) + ) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): StreamingScan.do_evaluate([scan], scan, context=ctx) + + +def test_split_scan_do_evaluate_missing_prefetch_metadata() -> None: + paths = ["/some/missing/file.parquet"] + parquet_options = ParquetOptions(prefetch_file_metadata=True) + context = IRExecutionContext() + schema = {"x": DataType(pl.Int64())} + + with pytest.raises( + AssertionError, + match=( + r"Parquet file metadata was not prefetched for paths: " + r"\['/some/missing/file\.parquet'\]\." + ), + ): + SplitScan.do_evaluate( + 0, + 4, + schema, + "parquet", + {}, + paths, + None, + 0, + -1, + None, + None, + None, + parquet_options, + context=context, + ) diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 19bede15d290..6dc45ccfa8ed 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -327,6 +327,7 @@ def test_parquet_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: m.setenv("CUDF_POLARS__PARQUET_OPTIONS__MAX_FOOTER_SAMPLES", "0") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__MAX_ROW_GROUP_SAMPLES", "0") m.setenv("CUDF_POLARS__PARQUET_OPTIONS__USE_RAPIDSMPF_NATIVE", "0") + m.setenv("CUDF_POLARS__PARQUET_OPTIONS__PREFETCH_FILE_METADATA", "1") # Test default engine = pl.GPUEngine() @@ -338,6 +339,7 @@ def test_parquet_options_from_env(monkeypatch: pytest.MonkeyPatch) -> None: assert config.parquet_options.max_footer_samples == 0 assert config.parquet_options.max_row_group_samples == 0 assert config.parquet_options.use_rapidsmpf_native is False + assert config.parquet_options.prefetch_file_metadata is True with monkeypatch.context() as m: m.setenv("CUDF_POLARS__PARQUET_OPTIONS__CHUNKED", "foo") @@ -416,6 +418,7 @@ def test_fallback_mode_default(monkeypatch: pytest.MonkeyPatch) -> None: "max_footer_samples", "max_row_group_samples", "use_rapidsmpf_native", + "prefetch_file_metadata", ], ) def test_validate_parquet_options(option: str) -> None: diff --git a/python/cudf_polars/tests/test_scan.py b/python/cudf_polars/tests/test_scan.py index 7a9a4f2bb10f..6fc63aed9de9 100644 --- a/python/cudf_polars/tests/test_scan.py +++ b/python/cudf_polars/tests/test_scan.py @@ -16,12 +16,15 @@ import polars as pl +from cudf_polars.containers import DataType +from cudf_polars.dsl.ir import IRExecutionContext, Scan from cudf_polars.testing.asserts import ( assert_gpu_result_equal, assert_ir_translation_raises, ) from cudf_polars.testing.engine_utils import is_streaming_engine from cudf_polars.testing.io import make_partitioned_source +from cudf_polars.utils.config import ParquetOptions from cudf_polars.utils.versions import ( POLARS_VERSION_LT_138, POLARS_VERSION_LT_139, @@ -168,6 +171,52 @@ def test_negative_slice_pushdown_raises(engine: pl.GPUEngine, tmp_path): assert_ir_translation_raises(q, engine, NotImplementedError) +@pytest.mark.parametrize("chunked", [False, True], ids=["single_read", "chunked"]) +def test_scan_parquet_prefetch_file_metadata( + tmp_path: Path, df: pl.DataFrame, *, chunked: bool +): + make_partitioned_source(df, tmp_path / "file", "parquet") + q = pl.scan_parquet(tmp_path / "file") + engine = pl.GPUEngine( + executor="in-memory", + raise_on_fail=True, + parquet_options={ + "chunked": chunked, + "prefetch_file_metadata": True, + }, + ) + assert_gpu_result_equal(q, engine=engine) + + +def test_scan_do_evaluate_missing_prefetch_metadata() -> None: + paths = ["/some/missing/file.parquet"] + parquet_options = ParquetOptions(prefetch_file_metadata=True) + context = IRExecutionContext() + schema = {"a": DataType(pl.Int64())} + + with pytest.raises( + AssertionError, + match=( + r"Parquet file metadata was not prefetched for paths: " + r"\['/some/missing/file\.parquet'\]\." + ), + ): + Scan.do_evaluate( + schema, + "parquet", + {}, + paths, + None, + 0, + -1, + None, + None, + None, + parquet_options, + context=context, + ) + + def test_scan_unsupported_raises(engine: pl.GPUEngine, tmp_path): df = pl.DataFrame({"a": [1, 2, 3]}) diff --git a/python/cudf_polars/tests/test_select.py b/python/cudf_polars/tests/test_select.py index f37c2d195d92..b5ab87d26042 100644 --- a/python/cudf_polars/tests/test_select.py +++ b/python/cudf_polars/tests/test_select.py @@ -8,10 +8,12 @@ import polars as pl +from cudf_polars.dsl.ir import IRExecutionContext, Scan from cudf_polars.testing.asserts import ( assert_gpu_result_equal, assert_ir_translation_raises, ) +from cudf_polars.utils.config import ParquetOptions def test_select(engine: pl.GPUEngine): @@ -147,3 +149,84 @@ def test_select_fast_count_parquet_skip_rows( q = pl.scan_parquet(file).slice(1, 5).select(pl.len()) assert_gpu_result_equal(q, engine=engine) + + +PARQUET_FAST_COUNT_ROWS = 10 + + +@pytest.fixture(scope="module") +def parquet_fast_count_df() -> pl.DataFrame: + return pl.DataFrame({"a": range(PARQUET_FAST_COUNT_ROWS)}) + + +@pytest.fixture +def prefetch_engine() -> pl.GPUEngine: + return pl.GPUEngine( + executor="in-memory", + raise_on_fail=True, + parquet_options={"prefetch_file_metadata": True}, + ) + + +@pytest.fixture( + params=[ + pytest.param({"skip_rows": 0, "n_rows": None}, id="all_rows"), + pytest.param({"skip_rows": 3, "n_rows": None}, id="skip_rows"), + pytest.param({"skip_rows": 2, "n_rows": 4}, id="skip_rows_and_limit"), + pytest.param({"skip_rows": 0, "n_rows": 5}, id="n_rows"), + pytest.param({"skip_rows": 8, "n_rows": 10}, id="skip_near_end"), + pytest.param( + {"skip_rows": PARQUET_FAST_COUNT_ROWS, "n_rows": None}, + id="skip_all", + ), + ], +) +def parquet_scan_row_bounds(request) -> dict[str, int | None]: + return request.param + + +def test_select_fast_count_parquet_prefetch_metadata( + tmp_path, + parquet_fast_count_df: pl.DataFrame, + prefetch_engine: pl.GPUEngine, + parquet_scan_row_bounds: dict[str, int | None], +) -> None: + skip_rows = parquet_scan_row_bounds["skip_rows"] + assert skip_rows is not None + n_rows = parquet_scan_row_bounds["n_rows"] + + file = tmp_path / "data.parquet" + parquet_fast_count_df.write_parquet(file) + + if skip_rows == 0 and n_rows is None: + q = pl.scan_parquet(file) + elif skip_rows == 0: + q = pl.scan_parquet(file, n_rows=n_rows) + elif n_rows is None: + q = pl.scan_parquet(file).slice(skip_rows) + else: + q = pl.scan_parquet(file).slice(skip_rows, n_rows) + + q = q.select(pl.len()) + assert_gpu_result_equal(q, engine=prefetch_engine) + + +def test_get_parquet_row_count_from_metadata_missing_prefetch() -> None: + paths = ["/some/missing/file.parquet"] + parquet_options = ParquetOptions(prefetch_file_metadata=True) + context = IRExecutionContext() + + with pytest.raises( + AssertionError, + match=( + r"Parquet file metadata was not prefetched for paths: " + r"\['/some/missing/file\.parquet'\]\." + ), + ): + Scan._get_parquet_row_count_from_metadata( + paths, + skip_rows=0, + n_rows=-1, + parquet_options=parquet_options, + context=context, + ) From 8402d9f16b20e0b3592b16af4d4834b39b7baac1 Mon Sep 17 00:00:00 2001 From: Tom Augspurger Date: Tue, 9 Jun 2026 13:31:11 -0700 Subject: [PATCH 4/4] Followup to tom/cudf-sourceinfo-size --- python/cudf_polars/cudf_polars/dsl/ir.py | 131 +++++++++++++++--- .../cudf_polars/cudf_polars/streaming/io.py | 3 +- 2 files changed, 114 insertions(+), 20 deletions(-) diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 06058867eaa7..4fd79e248b36 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -114,6 +114,14 @@ ] +@dataclass(frozen=True) +class CachedParquetInfo: + """Metadata for a parquet file.""" + + source_info: plc.io.SourceInfo + metadata: list[plc.io.parquet_metadata.FileMetaData] + + @dataclass(frozen=True) class IRExecutionContext: """ @@ -133,16 +141,17 @@ class IRExecutionContext: Identifier for the query being executed. parquet_file_metadata A cache of parquet file metadata. The keys are the ``paths`` of Scan nodes - with a ``parquet`` type. The values are a list of ``FileMetaData`` objects + with a ``parquet`` type. The values are `CachedParquetInfo` objects storing + a ``SourceInfo`` with known size and a list of ``FileMetaData`` objects associated with those ``paths``. """ py_executor: concurrent.futures.ThreadPoolExecutor | None = field(default=None) get_cuda_stream: Callable[[], Stream] = field(default=get_cuda_stream) query_id: uuid.UUID = field(default_factory=uuid.uuid4) - parquet_file_metadata: dict[ - tuple[str, ...], list[plc.io.parquet_metadata.FileMetaData] - ] = field(default_factory=dict) + parquet_file_metadata: dict[tuple[str, ...], CachedParquetInfo] = field( + default_factory=dict + ) async def to_thread( self, func: Callable[P, T], /, *args: P.args, **kwargs: P.kwargs @@ -199,7 +208,10 @@ def stream_ordered_after(self, *dfs: DataFrame) -> Generator[Stream, None, None] @nvtx_annotate_cudf_polars(message="fetch_parquet_footers_for_paths") def _prefetch_parquet_footers_for_paths( paths: tuple[str, ...], -) -> tuple[tuple[str, ...], list[plc.io.parquet_metadata.FileMetaData]]: + filepath_sources: list[plc.io.types.FilepathSource], +) -> tuple[ + tuple[str, ...], plc.io.SourceInfo, list[plc.io.parquet_metadata.FileMetaData] +]: """ Prefetch parquet footers for a list of paths. @@ -210,6 +222,8 @@ def _prefetch_parquet_footers_for_paths( ---------- paths The tuple of paths to prefetch. These correspond to ``paths`` in a ``Scan`` node. + filepath_sources + The list of ``FilepathSource`` objects for the ``paths``. Returns ------- @@ -219,10 +233,11 @@ def _prefetch_parquet_footers_for_paths( metadata The list of ``FileMetaData`` objects for the ``paths``. """ - metadata = plc.io.parquet_metadata.read_parquet_footers( - plc.io.SourceInfo(list(paths)) - ) - return paths, metadata + # TODO: https://github.com/rapidsai/cudf/issues/22734, use object metadata from polars + # For now, we'll just use kvikio to explicitly get the size. + source_info = plc.io.SourceInfo(filepath_sources) + metadata = plc.io.parquet_metadata.read_parquet_footers(source_info) + return paths, source_info, metadata @nvtx_annotate_cudf_polars(message="prefetch_parquet_file_metadata_for_ir") @@ -244,7 +259,7 @@ def prefetch_parquet_file_metadata_for_ir( """ from cudf_polars.streaming.io import SplitScan, StreamingScan - groups = set() + groups: set[tuple[str, ...]] = set() for node in traversal([root]): if isinstance(node, StreamingScan): for scan in node.scans: @@ -270,16 +285,90 @@ def prefetch_parquet_file_metadata_for_ir( # We didn't create the executor, so we don't close it. cm = contextlib.nullcontext() + import kvikio + + def make_filepath_source( + path: str, paths: tuple[str, ...], index: int + ) -> tuple[plc.io.types.FilepathSource, tuple[str, ...], int]: + if path.startswith("s3://"): + with kvikio.RemoteFile.open_s3_url(path) as remote_file: + return ( + plc.io.types.FilepathSource(path, size=remote_file.nbytes()), + paths, + index, + ) + else: + return plc.io.types.FilepathSource(path), paths, index + + # We have a `list[Node]`, and each node has a `list[Path]` with individual + # paths. + # + # We have two kinds of work to do: + # + # 1. Get the size of each individual `path`, return a + # `plc.types.FilepathSource`. 2. Call `read_parquet_footers` on the + # `list[plc.types.FilepathSource]`. + # + # The second type of work can only *start* once *all* the `FilepathSource`s + # are available *for that node*. + # + # Also, *order of the `FilepathSource`s matters*. We can process the Nodes + # in any order, and we can do the actual size fetch / `FilepathSource` + # creation in any order. But the `list[FilepathSource]` we pass into + # `read_parquet_footers` *must* be in the same order as the `list[Path]` for + # the Node. + # + # Finally, we want to use as much concurrency as possible. + sources_by_node: dict[ + tuple[str, ...], list[plc.io.types.FilepathSource | None] + ] = {} + for paths in groups: + sources_by_node[paths] = [None] * len(paths) + # we probably also want some kind of counter, but w/e + + ordered_paths = [] + for paths in groups: + for i, path in enumerate(paths): + ordered_paths.append((path, paths, i)) + if missing_paths: with cm: - futures = [ - executor.submit(_prefetch_parquet_footers_for_paths, paths) - for paths in missing_paths + filepath_futures = [ + executor.submit(make_filepath_source, path, paths, index) + for path, paths, index in ordered_paths ] - for future in concurrent.futures.as_completed(futures): - paths, metadata = future.result() - context.parquet_file_metadata.setdefault(paths, metadata) + for future in concurrent.futures.as_completed(filepath_futures): + filepath_source, paths, index = future.result() + sources_by_node[paths][index] = filepath_source + + # TODO: this is where you can check to see if we're done + # But for now, just do it at the end after the for loop completes I guess. + + # Now we know that we have a FilepathSource for each path in the node. + # Submit the read_parquet_footers for each node. + # for paths, sources in sources_by_node.items(): + futures: list[ + concurrent.futures.Future[ + tuple[ + tuple[str, ...], + plc.io.SourceInfo, + list[plc.io.parquet_metadata.FileMetaData], + ] + ] + ] = [] + for paths, sources in sources_by_node.items(): + # assert all(sources) + futures.append( + executor.submit(_prefetch_parquet_footers_for_paths, paths, sources) + ) + + for future_ in concurrent.futures.as_completed(futures): + paths_, source_info, metadata = future_.result() + context.parquet_file_metadata.setdefault( + paths_, + CachedParquetInfo(source_info=source_info, metadata=metadata), + ) _BINOPS = { @@ -756,7 +845,8 @@ def _get_parquet_row_count_from_metadata( # pylibcudf. See https://github.com/rapidsai/cudf/issues/21428 if parquet_options.prefetch_file_metadata and context is not None: try: - parquet_metadatas = context.parquet_file_metadata[tuple(paths)] + cached_parquet_info = context.parquet_file_metadata[tuple(paths)] + parquet_metadatas = cached_parquet_info.metadata except KeyError as e: msg = ( f"Parquet file metadata was not prefetched for paths: {list(paths)}." @@ -909,7 +999,9 @@ def read_csv_header( elif typ == "parquet": if parquet_options.prefetch_file_metadata: try: - parquet_metadatas = context.parquet_file_metadata[tuple(paths)] + cached_parquet_info = context.parquet_file_metadata[tuple(paths)] + source_info = cached_parquet_info.source_info + parquet_metadatas = cached_parquet_info.metadata except KeyError as e: msg = ( f"Parquet file metadata was not prefetched for paths: {list(paths)}." @@ -919,6 +1011,7 @@ def read_csv_header( raise AssertionError(msg) from e else: parquet_metadatas = None + source_info = plc.io.SourceInfo(paths) filters = None if predicate is not None and row_index is None: @@ -930,7 +1023,7 @@ def read_csv_header( stream=stream, ) parquet_reader_options = ( - plc.io.parquet.ParquetReaderOptions.builder(plc.io.SourceInfo(paths)) + plc.io.parquet.ParquetReaderOptions.builder(source_info) .decimal_width(plc.TypeId.DECIMAL128) .build() ) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 34bcf34bac01..5143a75818a9 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -308,7 +308,8 @@ def do_evaluate( if parquet_options.prefetch_file_metadata: try: - parquet_metadatas = context.parquet_file_metadata[tuple(paths)] + cached_parquet_info = context.parquet_file_metadata[tuple(paths)] + parquet_metadatas = cached_parquet_info.metadata except KeyError as e: msg = ( f"Parquet file metadata was not prefetched for paths: {list(paths)}."