Improve parquet footer reading performance by speculatively reading footer bytes - #22782
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
github-actions
Bot
added
libcudf
TomAugspurger
added
improvement
|
/ok to test f6e7eca |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR implements speculative Parquet footer reading by introducing a configurable ChangesParquet metadata size hint and speculative footer reading
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
cpp/tests/io/parquet_reader_test.cpp (1)
2972-3004: ⚡ Quick winHIGH: Add explicit
metadata_size_hint == 0boundary case in call-count testsLine 2998 and Line 3071 currently validate only
full_hintandminimal_hint=8. Please also asserthint=0behavior so the “disable speculative read” contract is regression-tested in both metadata and footer paths.Suggested test addition
TEST_F(ParquetMetadataReaderTest, MetadataSizeHintReadCallCount) { @@ EXPECT_EQ(tracking_source.reads()[0].second, minimal_hint); + + tracking_source.reset(); + constexpr size_t zero_hint = 0; + auto const metadata_zero_hint = cudf::io::read_parquet_metadata(source, zero_hint); + EXPECT_EQ(metadata_zero_hint.num_rows(), num_rows); + ASSERT_EQ(tracking_source.reads().size(), 2); + EXPECT_EQ(tracking_source.reads()[0].first, tracking_source.size() - 8); + EXPECT_EQ(tracking_source.reads()[0].second, 8); } @@ TEST_F(ParquetMetadataReaderTest, FooterSizeHintReadCallCount) { @@ EXPECT_EQ(tracking_source.reads()[0].second, minimal_hint); + + tracking_source.reset(); + constexpr size_t zero_hint = 0; + auto const zero_hint_footers = cudf::io::read_parquet_footers(sources, zero_hint); + ASSERT_EQ(zero_hint_footers.size(), 1); + EXPECT_EQ(zero_hint_footers[0].num_rows, num_rows); + ASSERT_EQ(tracking_source.reads().size(), 2); + EXPECT_EQ(tracking_source.reads()[0].first, tracking_source.size() - 8); + EXPECT_EQ(tracking_source.reads()[0].second, 8); }As per coding guidelines: “Test functions must cover edge cases: empty input, null values, sliced columns, boundary sizes, multi-block sizes”.
Also applies to: 3045-3078
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/io/parquet_reader_test.cpp` around lines 2972 - 3004, Add a third subcase calling cudf::io::read_parquet_metadata(source, 0) inside ParquetMetadataReaderTest::MetadataSizeHintReadCallCount and the other similar test (around lines 3045-3078) to explicitly exercise the metadata_size_hint == 0 boundary; verify the TrackingFooterDatasource::reads() shows that speculative suffix reads are disabled (i.e., the implementation issues a single precise footer read rather than a prior minimal speculative read) by asserting the reads count and that the first read offset/length correspond to the exact footer fetch rather than a minimal suffix read. Ensure you reference the test name MetadataSizeHintReadCallCount and the TrackingFooterDatasource instance (tracking_source) when adding the assertions.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/pylibcudf/tests/io/test_parquet.py`:
- Around line 295-353: Add edge-case tests for metadata_size_hint: extend the
existing tests (e.g., test_read_parquet_metadata_with_size_hint,
test_read_parquet_footers_with_size_hint, or add new functions) to call
plc.io.parquet_metadata.read_parquet_metadata and
plc.io.parquet_metadata.read_parquet_footers with metadata_size_hint=0 and
assert correct behavior (should succeed and produce same metadata as default),
and add a test that passes a negative metadata_size_hint (e.g., -1) to those
same functions and asserts they raise ValueError; reference SourceInfo,
plc.io.parquet_metadata.read_parquet_metadata,
plc.io.parquet_metadata.read_parquet_footers, and plc.io.parquet.read_parquet
(if using pre-materialized footers) when adding these assertions.
---
Nitpick comments:
In `@cpp/tests/io/parquet_reader_test.cpp`:
- Around line 2972-3004: Add a third subcase calling
cudf::io::read_parquet_metadata(source, 0) inside
ParquetMetadataReaderTest::MetadataSizeHintReadCallCount and the other similar
test (around lines 3045-3078) to explicitly exercise the metadata_size_hint == 0
boundary; verify the TrackingFooterDatasource::reads() shows that speculative
suffix reads are disabled (i.e., the implementation issues a single precise
footer read rather than a prior minimal speculative read) by asserting the reads
count and that the first read offset/length correspond to the exact footer fetch
rather than a minimal suffix read. Ensure you reference the test name
MetadataSizeHintReadCallCount and the TrackingFooterDatasource instance
(tracking_source) when adding the assertions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9c836288-7803-42d7-87b3-8050db6734d2
📒 Files selected for processing (16)
cpp/include/cudf/io/config_utils.hppcpp/include/cudf/io/detail/parquet.hppcpp/include/cudf/io/parquet_io_utils.hppcpp/include/cudf/io/parquet_metadata.hppcpp/src/io/functions.cppcpp/src/io/parquet/io_utils/parquet_io_utils.cppcpp/src/io/parquet/reader_impl.cppcpp/src/io/parquet/reader_impl_helpers.cppcpp/src/io/parquet/reader_impl_helpers.hppcpp/src/io/utilities/config_utils.cppcpp/tests/io/parquet_reader_test.cpppython/pylibcudf/pylibcudf/io/parquet_metadata.pxdpython/pylibcudf/pylibcudf/io/parquet_metadata.pyipython/pylibcudf/pylibcudf/io/parquet_metadata.pyxpython/pylibcudf/pylibcudf/libcudf/io/parquet_metadata.pxdpython/pylibcudf/tests/io/test_parquet.py
|
Some performance analysis: Detailsimport boto3
import nvtx
import pylibcudf as plc
bucket = "rapids-tpch"
key = "scale-10/nation/part.0.parquet"
url = f"s3://{bucket}/{key}"
sources = plc.io.types.SourceInfo([url])
with nvtx.annotate("warmup"):
plc.io.parquet_metadata.read_parquet_metadata(sources)
with nvtx.annotate("read_parquet_metadata"):
with nvtx.annotate("speculative"):
_ = plc.io.parquet_metadata.read_parquet_metadata(sources)
with nvtx.annotate("non-speculative"):
_ = plc.io.parquet_metadata.read_parquet_metadata(sources, metadata_size_hint=0)
with nvtx.annotate("read_parquet_footers"):
with nvtx.annotate("speculative"):
footers = plc.io.parquet_metadata.read_parquet_footers(sources)
with nvtx.annotate("non-speculative"):
footers_with_size = plc.io.parquet_metadata.read_parquet_footers(sources, metadata_size_hint=0)And here's a screnshot of the nsys profile captured with 
The most interesting bit is the It's not the easiest to follow since it's spread across threads, but this comes down to making one read (HTTP request) rather than two. Speculative read: 
Non-speculative read: 
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/tests/io/parquet_reader_test.cpp`:
- Around line 3254-3260: The test uses a fixed hint 64 which can be larger than
the actual footer start and make the speculative-read check flaky; instead
compute the hint from the actual footer length so the speculative read does NOT
include the full footer bytes (e.g., set hint = footer_len +
sizeof(file_ender_s) - 1) and pass that hint to read_parquet_metadata(source,
hint). Update the TailShortReadDatasource/parquet_bytes-based test setup to
derive footer_len from the generated parquet_bytes and use that derived hint in
the read_parquet_metadata call so the "Speculative metadata read did not include
full footer bytes" branch is reliably exercised.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 27f85695-3cdd-4740-8a5a-fbb90cbb475f
📒 Files selected for processing (3)
cpp/src/io/parquet/io_utils/parquet_io_utils.cppcpp/tests/io/parquet_reader_test.cpppython/pylibcudf/tests/io/test_parquet.py
🚧 Files skipped from review as they are similar to previous changes (2)
- python/pylibcudf/tests/io/test_parquet.py
- cpp/src/io/parquet/io_utils/parquet_io_utils.cpp
|
I would heavily cut down the added tests (see comments above). |
vuule
left a comment
vuule
left a comment
There was a problem hiding this comment.
Looks good, assuming memmove comments are resolved.
Thank you for iterating on this @TomAugspurger !
| 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); |
There was a problem hiding this comment.
Saves us one memcpy if we allocate footer_bytes before hand and directly host_read into it.
| 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); | |
| std::vector<uint8_t> footer_bytes(ender->footer_len); | |
| auto const missing_prefix_size = speculative_read_offset - footer_offset; | |
| datasource.host_read(footer_offset, missing_prefix_size, footer_bytes.data()); |
There was a problem hiding this comment.
Thanks @mhaseeb123. With that change, one of the (new) tests fails:
$ ./cpp/build/conda/cuda-12.9/release/gtests/PARQUET_TEST --gtest_filter=ParquetMetadataSizeHintTest.ReadParquet
Note: Google Test filter = ParquetMetadataSizeHintTest.ReadParquet
[==========] Running 1 test from 1 test suite.
[----------] Global test environment set-up.
[----------] 1 test from ParquetMetadataSizeHintTest
[ RUN ] ParquetMetadataSizeHintTest.ReadParquet
unknown file: Failure
C++ exception with description "CUDF failure at: /home/coder/cudf/cpp/src/io/parquet/reader_impl_helpers.cpp:324: Cannot initialize schema" thrown in the test body.
[ FAILED ] ParquetMetadataSizeHintTest.ReadParquet (114 ms)
[----------] 1 test from ParquetMetadataSizeHintTest (114 ms total)
[----------] Global test environment tear-down
[==========] 1 test from 1 test suite ran. (114 ms total)
[ PASSED ] 0 tests.
[ FAILED ] 1 test, listed below:
[ FAILED ] ParquetMetadataSizeHintTest.ReadParquet
And when running all the tests with ./cpp/build/conda/cuda-12.9/release/gtests/PARQUET_TEST there are a few more, seemingly the ones that have large footers:
[----------] Global test environment tear-down
[==========] 455 tests from 129 test suites ran. (117895 ms total)
[ PASSED ] 451 tests.
[ FAILED ] 4 tests, listed below:
[ FAILED ] ParquetReaderTest.TableTooLargeOverflows
[ FAILED ] ParquetChunkedWriterTest.LargeTables
[ FAILED ] ParquetChunkedWriterTest.ManyTables
[ FAILED ] ParquetMetadataSizeHintTest.ReadParquet
It seems to me like
datasource.host_read(footer_offset, missing_prefix_size, footer_bytes.data());doesn't get us the full footer. IIUC, footer_bytes is a newly allocated vec with the size of the full footer. After this host_read, the only initialized bytes will be [footer_start, missing_prefix), which will be a subset of the footer; just the bytes we failed to read in the speculative read. For example, if we have
- Footer size: 6405 Bytes
- Speculative read: 6400 Bytes
Then we'd have missing_prefix = 5 and footer_bytes would contain just the 5 bytes at the start of the footer plus uninitialized memory.
But I could easily be misreading things!
|
Here are some numbers from the benchmark Recall the tradeoff here: we're potentially reading too many bytes, in the hope that we capture the full footer. The possible outcomes are
https://gist.github.com/TomAugspurger/069abda3bbfe7728b910187a09ff7e5f is a little script to the output of this branch with speculative reading on and disabled (full results below). There are two cases where the time with speculative reading on is outside the noise band of the run with it off. The num_cols=64, num_row_groups=10 is faster. num_cols=512, num_row_groups=50 is slower. First, with the new default to speculatively read 64 KB: parquet_read_footer[0] Tesla V100-SXM2-32GB
Second, this branch with speculative reading disabled through the env var: parquet_read_footer[0] Tesla V100-SXM2-32GB
Also, if you're curious, only the first two cases (num_cols=64, num_row_groups=10) actually fit within the speculative read size. |
|
And comparing to So that one benchmark where we do the one-shot footer read is faster. The benchmark with the largest footer does appear to be slower: we fail to read the footer in a single read, and pay the cost of the memcpys. |
I feel like this is a good result. The only slower case has quite a few columns and row groups. In reality, such file would be huge and the few milliseconds lost reading the footer would not affect the total time significantly. |
Yep. As another point of reference, all of the footers are less that 64 KB for the the tpc-h files generated by tpch-rs at scale-1000. The largest is ~24 KB. |
|
/merge |
|
Thanks for the reviews! |
4 participants
Description
This PR aims to improve the performance of reading parquet footers (and all functions calling it, including
read_parquet) on high-latency storage systems like S3. On these storage systems, each read call translates to an HTTP range request, which takes 10s - 100s of ms to complete. Onmain, we make two read requests in order to only read exactly the bytes necessarynbytes required to read the actual footerOn high-latency storage systems, it's faster to make just one read (HTTP request), even if that read happens to be larger than is strictly necessary. It's (much) faster to read, say 64 KB in a single HTTP request than it is to read 8 bytes + (say) 32 KB in two HTTP requests.
This PR updates all the sites that eventually read the footer, including
read_parquet,read_parquet_metadata, etc. to accept ametadata_size_hint. If it's not provided, we fall back to a default controlled in our configuration system. The default is 64KB (matching arrow C++ and polars, smaller than DuckDB and datafusion).Speculative reading can be disabled at call sites by setting
metadata_size_hint=0, or globally by settingLIBCUDF_PARQUET_METADATA_SIZE_HINT=0. Either will restore the old behavior of reading exactly 8 bytes for the suffix, and then exactly the footer size.Python bindings in pylibcudf are added too.
Closes #22668
#22782 (comment) has some performance analysis when reading from S3.
Checklist