Write variable bit-width keys for Parquet dictionary encoded pages

[mhaseeb123]

Description

Contributes to #13995

This PR enables the Parquet writer to assign lower dictionary indices to elements appearing earlier in each column chunk thereby writing variable (reduced) number of dictionary key bits for earlier pages.

New algorithm:

1. map_insert_fn now inserts this pair in the static map: {row_idx, frag_idx} instead of {row_idx, row_idx}. frag_idx is just blockIdx.x and indicates which page fragment actually inserted this entry in the static map (CAS race dependent but earlier thread blocks generally win this race)
2. map_insert_fn also writes the PageFragment::num_dict_vals field for each page fragment to keep a record of number of unique values inserted by it.
3. Then in collect_map_entries_kernel (where we assign keys to dict values), we optionally assign spatially-local keys to dictionary values inserted by the same page fragment. See algorithm details here
4. Finally, we launch compute_page_dict_bits_kernel to compute the max bit width required to encode dict keys for each parquet page. See algorithm here

See follow up PR #22323 that deterministically writes the index of the first page fragment that sees each dictionary key

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6b7c490d-e7c7-4dcf-9cd0-4eed59918ad3

📥 Commits

Reviewing files that changed from the base of the PR and between ac9f3d3 and 1ceaaef.

📒 Files selected for processing (1)

• cpp/tests/io/parquet_writer_test.cpp

🚧 Files skipped from review as they are similar to previous changes (1)

• cpp/tests/io/parquet_writer_test.cpp

---

📝 Walkthrough

Summary by CodeRabbit

• New Features

  • Added a Parquet dictionary-encoding benchmark to measure write performance, memory, file-size, and per-page RLE bit-widths with configurable string distributions.

• Improvements

  • Compute and record dictionary RLE bit-widths at the page level for more accurate encoding and size estimation.
  • Writer/encoder updates to materialize page-level dictionary metadata for more reliable writes.

• Tests

  • Added/updated tests to validate per-page dictionary bit-widths and variable bit-width behavior.

Walkthrough

This PR extends the cuDF Parquet writer to compute per-page dictionary RLE bit widths after page boundaries are finalized, replacing chunk-level dictionary settings. It widens fragment structures, introduces GPU kernels for per-page bit-width calculation, and updates writer orchestration to call the new compute function with finalized pages. Tests validate the per-page behavior across variable-distribution scenarios.

Changes

Per-Page Dictionary RLE Bit-Width Computation

Layer / File(s)	Summary
Data Structures cpp/src/io/parquet/parquet_gpu.hpp	PageFragment::num_rows and num_dict_vals widened from uint16_t to size_type; EncPage gains uint8_t dict_rle_bits field.
GPU API Declarations cpp/src/io/parquet/parquet_gpu.cuh	populate_chunk_hash_maps and collect_map_entries signatures updated to include fragment span; new compute_per_page_dict_bits API declared.
GPU Dictionary Kernels cpp/src/io/parquet/chunk_dict.cu	Hash insertion/lookup functors refactored with fragment-aware indexing; new populate_chunk_hash_maps_kernel, collect_map_entries_kernel, and compute_per_page_dict_bits_kernel with 2D fragment handling and warp-level reduction.
Page Encoding cpp/src/io/parquet/page_enc.cu	init_frag_state bounds num_rows with cuda::std::min; gpuInitPages sets per-page dict_rle_bits; gpuEncodeDictPages uses page.dict_rle_bits instead of chunk-level default.
Writer Orchestration cpp/src/io/parquet/writer_impl.cu	build_chunk_dictionaries takes non-const fragments; num_fragments set during chunk initialization and remapping; compute_per_page_dict_bits called after page boundaries finalized.
Benchmarks & Tests cpp/benchmarks/CMakeLists.txt, cpp/benchmarks/io/parquet/parquet_writer_dict.cpp, cpp/tests/io/parquet_misc_test.cpp, cpp/tests/io/parquet_writer_test.cpp	New NVBench for dictionary encoding with tunable page distributions; updated DictionaryTest to validate max bit-width across all pages; new VariableBitWidthDictEncoding test with frequent/rare page regions.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 40.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title directly and concisely describes the main feature: enabling variable (reduced) dictionary key bit-widths for Parquet pages.
Description check	✅ Passed	The description clearly explains the optimization, outlines the new algorithm steps, references related issues/PRs, and provides context for the changes.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

• Generate code and open pull requests
• Plan features and break down work
• Investigate incidents and troubleshoot customer tickets together
• Automate recurring tasks and respond to alerts with triggers
• Summarize progress and report instantly

Built for teams:

• Shared memory across your entire org—no repeating context
• Per-thread sandboxes to safely plan and execute work
• Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 1

♻️ Duplicate comments (1)

cpp/src/io/parquet/chunk_dict.cu (1)

136-143: ⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

The fragment hint is still nondeterministic.

This records the frag_idx of whichever block wins the insert, not the earliest fragment that observed the value. On out-of-order block scheduling, the same input can still receive different dictionary IDs and therefore different per-page bit widths/file sizes across runs, so the new exact page-bit assertions remain flaky until first-fragment stamping is deterministic.

🤖 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/src/io/parquet/chunk_dict.cu` around lines 136 - 143, The current
insertion uses map_insert_ref.insert(slot_type{static_cast<key_type>(val_idx),
frag_idx}) which records whichever block wins the race, producing
nondeterministic frag_idx; change this to use the static_map's insert_or_apply
API with cuco::op::min so the stored value for key_type(val_idx) is always the
minimum frag_idx seen (i.e., first fragment). Locate the insert call in
chunk_dict.cu (map_insert_ref.insert and the surrounding uniq_elem_size lambda)
and replace the insert with insert_or_apply passing a min operation
(cuco::op::min) so ties resolve deterministically to the earliest fragment
index.

🤖 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/src/io/parquet/writer_impl.cu`:
- Around line 2125-2127: compute_per_page_dict_bits is being called using stale
fragment-local dictionary IDs assigned earlier from row_group_fragments, so when
the second pass mutates ck.fragments/ck.num_fragments the per-page bit widths no
longer match the actual fragments that formed the pages; to fix this, ensure
compute_per_page_dict_bits is invoked with page/fragment info reflecting the
post-pass fragment ordering — either recompute fragment-local dictionary IDs
(re-run the row_group_fragments assignment) or rebuild the pages buffer from the
updated ck.fragments/ck.num_fragments before calling compute_per_page_dict_bits,
and move this call to after any mutation of ck.fragments/ck.num_fragments so the
bit-widths are computed against the final fragment layout.

---

Duplicate comments:
In `@cpp/src/io/parquet/chunk_dict.cu`:
- Around line 136-143: The current insertion uses
map_insert_ref.insert(slot_type{static_cast<key_type>(val_idx), frag_idx}) which
records whichever block wins the race, producing nondeterministic frag_idx;
change this to use the static_map's insert_or_apply API with cuco::op::min so
the stored value for key_type(val_idx) is always the minimum frag_idx seen
(i.e., first fragment). Locate the insert call in chunk_dict.cu
(map_insert_ref.insert and the surrounding uniq_elem_size lambda) and replace
the insert with insert_or_apply passing a min operation (cuco::op::min) so ties
resolve deterministically to the earliest fragment index.

🪄 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: 49a19954-7286-4c84-bc55-adbbef540835

📥 Commits

Reviewing files that changed from the base of the PR and between c9ad1c5 and 7cc001c.

📒 Files selected for processing (9)

• cpp/benchmarks/CMakeLists.txt
• cpp/benchmarks/io/parquet/parquet_writer_dict.cpp
• cpp/src/io/parquet/chunk_dict.cu
• cpp/src/io/parquet/page_enc.cu
• cpp/src/io/parquet/parquet_gpu.cuh
• cpp/src/io/parquet/parquet_gpu.hpp
• cpp/src/io/parquet/writer_impl.cu
• cpp/tests/io/parquet_misc_test.cpp
• cpp/tests/io/parquet_writer_test.cpp

[mhaseeb123]

/ok to test 1ceaaef

[mhaseeb123]

/ok to test d2b1028

[mhaseeb123]

/ok to test e84a14f

[mhaseeb123]

/merge

[ttnghia]

CMake approval.