Fix Parquet V2 inputs to decompression scratch queries - #24006
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (4)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 SummarySummary by CodeRabbit
WalkthroughChangesParquet chunked-read handling
Estimated code review effort: 3 (Moderate) | ~25 minutes Merge Risk: ⚪ Minimal · up to Chunked Parquet reads now size Zstd decompression scratch space from actual compressed value payloads, avoiding V2 level bytes and skipped pages. The changed behavior is covered across relevant page and reader configurations, with no current merge-blocking risk identified. Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 8 functions across 3 files. (1 skipped: 1 unsupported.)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Comment |
| /** | ||
| * @brief Returns the compressed values passed to the decompressor, excluding V2 level bytes. | ||
| */ | ||
| struct get_decompression_input { |
There was a problem hiding this comment.
The definition of "what the decompressor actually consumes" now lives in three places: codec_stats::add_pages (reader_impl_chunking_utils.cu:121-146), set_parameters (reader_impl_chunking_utils.cu:529-575), and this new functor. Divergence between the first two and the scratch query is precisely the bug being fixed here, so leaving three copies invites a repeat.
Declaring the functor CUDF_HOST_DEVICE and calling it from both host-side loops would collapse the duplication and make the invariant enforced by construction rather than by comment.
| { | ||
| return {parquet_compression_support(chunks[p.chunk_idx].codec).first, | ||
| auto const codec = parquet_compression_support(chunks[p.chunk_idx].codec).first; | ||
| if (get_decompression_input{}(p).empty()) { return {codec, 0, 0, 0}; } |
There was a problem hiding this comment.
With get_decomp_info now excluding empty-payload V2 pages, the estimate can sit slightly below what decompression actually requests: codec_stats::add_pages still folds those pages' uncompressed_page_size into codec.total_decomp_size and max_decompressed_size, and both are passed straight to cudf::io::detail::decompress at reader_impl_chunking_utils.cu:618-619.
The gap is small (such pages carry only level bytes) and the previous code overestimated, so this is not a blocker, but it is the same class of divergence the PR sets out to remove and is worth either closing or calling out in the comment.
| ? get_decompression_input{}(page) | ||
| : device_span<uint8_t const>{}; | ||
| }); | ||
| // Copy only non-null spans |
There was a problem hiding this comment.
The copy_if predicate just below still keys off span.data() != nullptr. Now that empty spans are produced deliberately and carry meaning, not span.empty() expresses the intent directly and is robust if get_decompression_input ever returns a non-null zero-length span.
| check_chunked_read(filepath, expected, use_metadata); | ||
| } | ||
|
|
||
| INSTANTIATE_TEST_SUITE_P(PageVersionAndReader, |
There was a problem hiding this comment.
The 16 reader cases are more than this bug needs, and two of the parameterization axes cannot interact with the code being fixed.
use_metadata doubles every case to compare the two chunked_parquet_reader constructors, but both converge on the same chunking and scratch-estimation code once the footer is parsed. dictionary doubles it again, yet dictionary pages never carry PAGEINFO_FLAGS_V2, so is_compressed is always true and the level offset always zero for them, yielding no variation in the three conditions the fix turns on.
Dropping use_metadata and dictionary leaves V2 with chunk-level compression, V2 with page-level compression, and a V1 regression guard, which covers the defect at about a third of the runtime.
| tmp_env_var const nvcomp{nvcomp_policy_env_var, "ALWAYS"}; | ||
| tmp_env_var const host_decomp{host_decomp_env_var, "OFF"}; | ||
| auto const stream = cudf::get_default_stream(); | ||
| // A Zstd frame containing 32,768 repetitions of '*'. |
There was a problem hiding this comment.
The hand-written 19-byte Zstd frame is opaque and pins the test to a specific frame encoding. Producing it with cudf::io::detail::compress(compression_type::ZSTD, ...) at test setup would be self-describing and would keep the test valid if the accepted frame layout ever shifts.
| NVBENCH_BENCH_TYPES(BM_parquet_read_subrowgroup_chunks, NVBENCH_TYPE_AXES(d_type_list)) | ||
| .set_name("parquet_read_subrowgroup_chunks") | ||
| .add_string_axis("io_type", {"DEVICE_BUFFER"}) | ||
| .add_string_axis("compression_type", {"SNAPPY"}) |
There was a problem hiding this comment.
Adding compression_type, write_v2_headers, and page_level_compression as single-value axes is forced by sharing one callable across two registrations, but it changes every existing parquet_read_subrowgroup_chunks state string and so breaks continuity with previously recorded results. A thin wrapper function that hardcodes the Snappy/V1 values for the original benchmark would keep the old axis set intact.
mhaseeb123
left a comment
mhaseeb123
left a comment
There was a problem hiding this comment.
Same comments as @vuule plus the following
| } | ||
| }; | ||
|
|
||
| TEST_F(ParquetScratchTest, DecompressionInputs) |
There was a problem hiding this comment.
I think this and SkippedPagesDoNotReduceScratchEstimate tests should be removed and the fix should be covered via Parquet write/read. This also removes the internal-header include and the CUDF_EXPORT on compute_decompression_scratch_sizes util.
Note that pytests in pylibcudf or cudf-python can also be added if libcudf alone can't exercise the bug path.
Please move any generic decompression tests to the respective suite in tests/io/comp.
| .add_int64_axis("row_group_size_bytes", {0}) | ||
| .add_int64_axis("row_group_size_rows", {0}); | ||
|
|
||
| // Keep the existing 512 MiB matrix unchanged; cover V2 scratch queries with 12 small cases. |
There was a problem hiding this comment.
I think these extra benchmarks should be removed as we don't really need a benchmark for a bug fix (unless the fix should affect the performance). We should cover V2 scratch queries in tests.
3 participants
Description
Chunked Parquet reads with a non-zero input pass limit can hang or encounter an illegal device access on Zstd DataPageV2 files. The extended decompression scratch-size query receives each whole page, including uncompressed repetition/definition levels and pages whose values are not compressed. The decompressor already handles these cases when preparing its actual input buffers.
Use a shared input-selection helper to skip uncompressed V2 pages, advance past level bytes, and exclude pages without a remaining payload. Apply that selection to both cumulative and total decompression-size estimates, so skipped pages cannot dilute the extended/legacy scratch adjustment ratio. Retain full page output sizes as conservative bounds, consistent with the decompressor's codec statistics; prefixes with no compressed pages require no scratch query.
Add 16 reader cases covering V1, V2 with default chunk-level compression, V2 with page-level compression, nullable/list columns, dictionary/plain encoding, mixed compressed/uncompressed pages, and empty pages within Zstd chunks. Each checks values and nulls with small and large pass limits through direct file input and pre-parsed footers. Two direct regression tests check input pointers/lengths and verify that inserting skipped pages cannot reduce the scratch estimate for the same compressed pages.
Keep the existing
parquet_read_subrowgroup_chunksbenchmark's 144 Snappy/V1 configurations at 512 MiB. Add a separateparquet_read_v2_scratchbenchmark with only 12 Zstd/V2 configurations at a committed default of 8 MiB, covering integer/string/list inputs, page-level compression on/off, and pass limits of 0/500,000 bytes.Closes #24002.
Validation
PARQUET_TEST,HYBRID_SCAN_TEST, andPARQUET_READER_CHUNKS_NVBENCHfrom main ata8ad2045e17034209e29cec01a6b4a51d87764d7(26.10), using a separate worktree/build and an isolated dependency environment on NVIDIA GB10 / Linux aarch64: CUDA 13.3, GCC 14.4, RMM/KvikIO 26.10 nightlies, nvCOMP 5.3.0.16.PARQUET_TEST: 566 passed, 1 skipped, 4 disabled.HYBRID_SCAN_TEST: all 97 passed.compute-sanitizer --tool memcheck: all 18 new cases passed, with 0 errors.DecompressionInputsfail on V2 offsets and skipped pages. Restoring unfiltered page-size accounting makesSkippedPagesDoNotReduceScratchEstimatefail: its first compressed-page estimate falls from 3,279 to 1,588 bytes when skipped pages are inserted. These tests detect the defects without relying on nvCOMP hanging on invalid input. Restored the fix and rebuilt before running the full suites and sanitizer above.Checklist