From 786f1e02ad5de998a9aa8069feb0221807e1bfb6 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Wed, 22 Apr 2026 20:38:33 +0000 Subject: [PATCH 01/28] Phase 0 --- cpp/benchmarks/CMakeLists.txt | 5 +- .../io/parquet/parquet_writer_dict.cpp | 174 ++++++++++++++++++ 2 files changed, 178 insertions(+), 1 deletion(-) create mode 100644 cpp/benchmarks/io/parquet/parquet_writer_dict.cpp diff --git a/cpp/benchmarks/CMakeLists.txt b/cpp/benchmarks/CMakeLists.txt index 066716e85f25..585266c50a05 100644 --- a/cpp/benchmarks/CMakeLists.txt +++ b/cpp/benchmarks/CMakeLists.txt @@ -284,7 +284,10 @@ ConfigureNVBench(BITMASK_NVBENCH bitmask/bitmask_and.cpp bitmask/set_null_mask.c # ################################################################################################## # * parquet writer benchmark ---------------------------------------------------------------------- ConfigureNVBench( - PARQUET_WRITER_NVBENCH io/parquet/parquet_writer.cpp io/parquet/parquet_writer_chunks.cpp + PARQUET_WRITER_NVBENCH + io/parquet/parquet_writer.cpp + io/parquet/parquet_writer_chunks.cpp + io/parquet/parquet_writer_dict.cpp ) # ################################################################################################## diff --git a/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp b/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp new file mode 100644 index 000000000000..0e9214deb5ee --- /dev/null +++ b/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp @@ -0,0 +1,174 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace { +// Dict-encoding-focused Parquet writer benchmark. Built to isolate the two +// encoder pathologies reported in https://github.com/rapidsai/cudf/issues/13995: +// +// 1. `dict_id` assignment order. cuDF's `collect_map_entries_kernel` walks +// the hash map in slot order, so the `dict_id` assigned to a value is +// uncorrelated with how early it appears in the column. Early pages +// therefore end up referencing high-indexed entries, defeating any +// per-page bit-width savings. +// 2. Chunk-wide RLE bit width. `build_chunk_dictionaries` derives one +// `dict_rle_bits = ceil(log2(num_dict_entries))` for the entire chunk, +// so a page touching only a handful of entries still bit-packs at the +// chunk-wide width. +// +// Workload shape (deliberately the simplest construction that exposes both +// problems on a single row group / single chunk): +// +// * 1 INT64 column, 1 row group = 1 chunk, 10 pages per chunk. +// INT64 is the type reported in the issue; fixed width keeps the +// per-row byte count deterministic so the file-size delta is driven +// entirely by how the encoder packs the dict-index stream. +// * cardinality = 64,000 (ceil(log2) = 16 bits chunk-wide). +// * Pages 0..`hot_pages-1` (the "common" pages) draw uniformly from a +// small "frequent set" of `frequent_set_size` values shared across +// all common pages. +// * Pages `hot_pages..pages_per_chunk-1` (the "rare" pages) draw +// uniformly from the remaining `cardinality - frequent_set_size` +// values. These never appear in any common page, so only the rare +// pages force the chunk-wide bit width up. +// +// Under an ideal encoder (first-appearance ordering + per-page bit width) +// the common pages need only `ceil(log2(frequent_set_size))` bits per +// value while the rare pages need `ceil(log2(cardinality))` bits. With +// the constants below (frequent_set_size = 64), that is 6 bits vs. 16 +// bits, saving ~10 bits/value on every common-page row. The current +// cuDF encoder produces 16 bits on every page, so the headline +// `encoded_file_size` metric cleanly resolves the optimization target. + +// Distribution shape — see file header. These are intentionally fixed so that +// the encoded-file-size number is directly comparable across phases. Adjust +// together with the phase baseline if the shape changes. +constexpr cudf::size_type num_cols = 1; +constexpr cudf::size_type cardinality = 64'000; +constexpr cudf::size_type pages_per_chunk = 10; +constexpr cudf::size_type hot_pages = pages_per_chunk - 2; +constexpr cudf::size_type frequent_set_size = 64; +constexpr std::uint32_t dict_rng_seed = 0xC0DEFACE; + +// Build the row-to-dictionary-index mapping on host for a single column. +// Index space (disjoint, exactly covers [0, cardinality)): +// * frequent set: [0, frequent_set_size) +// * rare set: [frequent_set_size, cardinality) +template +std::vector build_numeric_column(cudf::size_type num_rows, cudf::size_type page_size_rows) +{ + CUDF_EXPECTS(num_rows == pages_per_chunk * page_size_rows, + "num_rows must equal pages_per_chunk * page_size_rows"); + CUDF_EXPECTS(frequent_set_size < cardinality, + "cardinality must leave room for a nonempty rare set"); + + std::mt19937 rng{dict_rng_seed}; + std::uniform_int_distribution freq_dist(0, frequent_set_size - 1); + std::uniform_int_distribution rare_dist(frequent_set_size, cardinality - 1); + auto const threshold = hot_pages * page_size_rows; + + cudf::size_type row_idx = 0; + std::vector values(num_rows); + std::generate_n(values.begin(), num_rows, [&]() { + return row_idx++ < threshold ? static_cast(freq_dist(rng)) : static_cast(rare_dist(rng)); + }); + + return values; +} + +[[nodiscard]] std::unique_ptr build_table(cudf::size_type num_rows, + cudf::size_type page_size_rows) +{ + auto const values = build_numeric_column(num_rows, page_size_rows); + std::vector> cols; + cols.reserve(num_cols); + cols.emplace_back( + cudf::test::fixed_width_column_wrapper(values.begin(), values.end()).release()); + return std::make_unique(std::move(cols)); +} + +} // namespace + +void BM_parq_write_dict_encoding(nvbench::state& state) +{ + auto const num_rows = static_cast(state.get_int64("num_rows")); + auto const page_size_rows = num_rows / pages_per_chunk; + + auto const tbl = build_table(num_rows, page_size_rows); + auto const view = tbl->view(); + + std::size_t encoded_file_size = 0; + + auto const mem_stats_logger = cudf::memory_stats_logger(); + state.set_cuda_stream(nvbench::make_cuda_stream_view(cudf::get_default_stream().value())); + state.exec(nvbench::exec_tag::timer | nvbench::exec_tag::sync, + [&](nvbench::launch&, auto& timer) { + cuio_source_sink_pair source_sink(io_type::FILEPATH); + + // Lift the byte caps well above the row caps so page boundaries + // are driven purely by `max_page_size_rows` (100K rows = 800KB + // for INT64, which exceeds the default 512KB page byte cap). + // This guarantees exactly `pages_per_chunk` pages per chunk, + // matching the host-side data layout. + constexpr std::size_t page_bytes_cap = std::size_t{64} << 20; + + timer.start(); + auto const write_opts = + cudf::io::parquet_writer_options::builder(source_sink.make_sink_info(), view) + .compression(cudf::io::compression_type::NONE) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .row_group_size_rows(num_rows) + .max_page_size_rows(page_size_rows) + .max_page_size_bytes(page_bytes_cap) + .build(); + cudf::io::write_parquet(write_opts); + timer.stop(); + + encoded_file_size = source_sink.size(); + }); + + state.add_element_count(static_cast(view.num_rows()), "rows"); + state.add_buffer_size( + mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage"); + state.add_buffer_size(encoded_file_size, "encoded_file_size", "encoded_file_size"); + + // TODO(phase 2): once per-page `dict_rle_bits` is plumbed into the writer + // path and exposed at the reader, replace this whole-file proxy with an + // exact min/max/mean across emitted pages. Until then we report aggregate + // bits-per-value, which tracks overall compression ratio but not the + // per-page distribution we ultimately want to shrink. + auto const total_values = static_cast(view.num_rows()) * num_cols; + if (total_values > 0) { + auto const bits_per_value_avg = + static_cast(encoded_file_size * 8) / static_cast(total_values); + state.add_element_count(bits_per_value_avg, "bits_per_value_avg"); + } +} + +NVBENCH_BENCH(BM_parq_write_dict_encoding) + .set_name("parquet_write_dict_encoding") + .set_min_samples(3) + .add_int64_axis("num_rows", {1'000'000}); From 08f02ae251c07ee88a40edad7aefebb5bd234694 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Wed, 22 Apr 2026 22:39:07 +0000 Subject: [PATCH 02/28] OPT: assign parquet dict_ids in first-appearance order (#13995) Repurpose the unused mapped_type slot in populate_chunk_hash_maps_kernel to carry the fragment index of the block that first inserts each value. Rewrite collect_map_entries_kernel to bucket dict_ids by that fragment index so pages that only reference earlier fragments' values see small max dict_index. No file-size delta in isolation; prerequisite for the per-page bit-width change landing next. Made-with: Cursor --- cpp/src/io/parquet/chunk_dict.cu | 162 +++++++++++++++++++++++++---- cpp/src/io/parquet/parquet_gpu.cuh | 11 ++ cpp/src/io/parquet/writer_impl.cu | 2 +- 3 files changed, 155 insertions(+), 20 deletions(-) diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index d24d82d518f1..55f3512e2f94 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -12,6 +12,7 @@ #include +#include #include #include @@ -19,7 +20,15 @@ namespace cudf::io::parquet::detail { namespace { constexpr int DEFAULT_BLOCK_SIZE = 256; -} + +// Upper bound on the number of fragments per column chunk that the +// shared-memory histogram in `collect_map_entries_kernel` can accommodate. +// A typical workload is 1M row groups / ~5000-row fragments ≈ 200 fragments +// per chunk, so 1024 is a comfortable ceiling. If a future workload exceeds +// this, `collect_map_entries_kernel` will trip its `cudf_assert` below and +// we should add a global-memory fallback rather than silently truncating. +constexpr size_type MAX_FRAGMENTS_PER_BLOCK = 1024; +} // namespace template struct equality_functor { @@ -93,7 +102,34 @@ struct map_insert_fn { // Insert tile_val_idx to hash map and count successful insertions. if (is_valid) { // Insert the keys using a single thread for best performance for now. - is_unique = map_insert_ref.insert(cuco::pair{val_idx, val_idx}); + // Stamp `slot.second` with the (approximate) lowest column-relative fragment + // index that inserted this value. `cuco::static_map_ref::insert` writes the + // value exactly once per key -- on the winning insertion -- so after this + // kernel, every non-empty slot's `second` field equals the `blockIdx.x` of + // the block that first won that slot. Block scheduling on NVIDIA GPUs is + // approximately monotonic in `blockIdx`, giving a strong first-appearance + // signal that `collect_map_entries_kernel` converts into monotone dict_ids. + // + // NOTE: `blockIdx.x` here is the *column-relative* fragment index (the + // populate grid is (num_fragments_per_col, num_cols)). Each `EncColumnChunk` + // covers a contiguous range of these indices within its column, so the + // collect kernel subtracts the chunk's first column-relative fragment to get + // a chunk-local histogram bucket. + // + // The pair is constructed as `slot_type` (exactly `value_type` for the map, + // 8 bytes total) rather than the tempting `cuco::pair{val_idx, hint}`. The + // latter would deduce `cuco::pair` because `val_idx` is + // `thread_index_type` (int64_t), and cuco's `packed_cas` path (selected when + // `sizeof(value_type) <= 8`) reinterprets the input as a single `uint64_t` + // -- it would then CAS the low 8 bytes (the `int64_t` key) into the slot and + // silently drop the payload, leaving `slot.second = 0`. + // + // TODO: when the cuco pin exposes `static_map::insert_or_apply` with + // `cuco::op::min`, switch to that for exact first-fragment semantics + // (PROBLEM.md §5.2 Option II). + auto const fragment_hint = static_cast(blockIdx.x); + is_unique = map_insert_ref.insert( + slot_type{static_cast(val_idx), fragment_hint}); uniq_elem_size = [&]() -> size_type { if (not is_unique) { return 0; } switch (col->physical_type) { @@ -226,32 +262,119 @@ CUDF_KERNEL void __launch_bounds__(block_size) end_value_idx); } +// Assigns monotone `dict_id`s bucketed by the fragment-hint that +// `populate_chunk_hash_maps_kernel` stamped into every slot's `second` field. +// After this kernel, `slot.second` has been overwritten with the final +// `dict_id` (an index into `chunk.dict_data`), and all values first inserted +// by fragment `f_i` are assigned `dict_id`s strictly less than values first +// inserted by fragment `f_j` for `f_i < f_j`. Per-page max `dict_index` is +// therefore monotone in page order whenever the underlying value distribution +// is, which is the invariant Phase 2 exploits for per-page bit-width savings. +// +// Algorithm (one block per chunk, three passes over `dict_map_size` slots): +// 1. Histogram: count non-empty slots by (fragment_hint - f_start). +// 2. BlockScan the histogram to exclusive-prefix offsets. +// 3. Claim `dict_id = atomicAdd(&fragment_cursor[bucket], 1)` per slot, +// overwrite slot.second, and write `dict_data[dict_id] = key`. +// +// Three passes of `dict_map_size / block_size` iterations are cheaper than +// the alternative of launching a separate `cub::DeviceHistogram` + +// `DeviceScan` per chunk (kernel-launch overhead dominates for many small +// chunks; see PROBLEM.md §5.1 and PHASE_1_ORDERING.md §1.2.2). template CUDF_KERNEL void __launch_bounds__(block_size) collect_map_entries_kernel(device_span const map_storage, - device_span chunks) + device_span chunks, + cudf::detail::device_2dspan frags) { + static_assert(block_size >= MAX_FRAGMENTS_PER_BLOCK, + "block_size must be >= MAX_FRAGMENTS_PER_BLOCK so one BlockScan thread backs " + "each histogram bucket."); + auto& chunk = chunks[blockIdx.x]; if (not chunk.use_dictionary) { return; } - auto t = threadIdx.x; - __shared__ cuda::atomic counter; - using cuda::std::memory_order_relaxed; - if (t == 0) { new (&counter) cuda::atomic{0}; } + auto const t = threadIdx.x; + auto const col_idx = chunk.col_desc_id; + auto const col_frag = frags[col_idx]; + + // Resolve the chunk's column-relative fragment range [f_start, f_start + n_frags). + // `chunk.fragments` points into `col_frag` (set in writer_impl.cu during + // row_group_fragments setup, before build_chunk_dictionaries runs), so the + // difference is the first fragment index stamped into any of this chunk's + // slots by populate_chunk_hash_maps_kernel. + __shared__ size_type f_start; + __shared__ size_type n_frags; + if (t == 0) { + f_start = static_cast(chunk.fragments - col_frag.data()); + size_type n = 0; + auto const total_col_frags = static_cast(col_frag.size()); + while (f_start + n < total_col_frags && col_frag[f_start + n].chunk == &chunk) { + ++n; + } + n_frags = n; + } + __syncthreads(); + + cudf_assert(n_frags > 0 && n_frags <= MAX_FRAGMENTS_PER_BLOCK && + "Fragments-per-chunk out of range for shared-memory histogram; raise " + "MAX_FRAGMENTS_PER_BLOCK or add a global-memory fallback."); + + // `fragment_count` starts as the per-bucket histogram, becomes per-bucket + // exclusive offsets after the scan, then is preserved as a reference for + // the MAX_DICT_SIZE overflow assert in Pass 3. `fragment_cursor` mirrors + // the offsets and is what threads `atomicAdd` into to claim dict_ids. + __shared__ size_type fragment_count[MAX_FRAGMENTS_PER_BLOCK]; + __shared__ size_type fragment_cursor[MAX_FRAGMENTS_PER_BLOCK]; + + using block_scan = cub::BlockScan; + __shared__ typename block_scan::TempStorage scan_storage; + + // Zero the histogram (only the live range; the tail past n_frags is + // untouched and never read). + if (t < n_frags) { fragment_count[t] = 0; } + __syncthreads(); + + auto* const chunk_slots = map_storage.data() + chunk.dict_map_offset; + auto const dict_map_size = chunk.dict_map_size; + + // Pass 1: histogram slot.second values into per-fragment buckets. + // `slot.second` is the column-relative fragment index that won the insert in + // `populate_chunk_hash_maps_kernel`; subtracting `f_start` normalizes it into + // a chunk-local bucket index `[0, n_frags)`. + for (size_type slot_idx = t; slot_idx < dict_map_size; slot_idx += block_size) { + auto const slot_key = chunk_slots[slot_idx].first; + if (slot_key != KEY_SENTINEL) { + auto const frag_local = chunk_slots[slot_idx].second - f_start; + cudf_assert(frag_local >= 0 && frag_local < n_frags && + "populate stamped a fragment hint outside this chunk's fragment range"); + atomicAdd(&fragment_count[frag_local], 1); + } + } + __syncthreads(); + + // Pass 2: in-block exclusive scan of the histogram -> offsets -> cursors. + { + size_type const per_thread_count = (t < n_frags) ? fragment_count[t] : 0; + size_type per_thread_offset = 0; + block_scan(scan_storage).ExclusiveSum(per_thread_count, per_thread_offset); + if (t < n_frags) { + fragment_count[t] = per_thread_offset; + fragment_cursor[t] = per_thread_offset; + } + } __syncthreads(); - // Iterate over all slots in the map. - for (; t < chunk.dict_map_size; t += block_size) { - auto* slot = map_storage.data() + chunk.dict_map_offset + t; - auto const key = slot->first; - if (key != KEY_SENTINEL) { - auto const loc = counter.fetch_add(1, memory_order_relaxed); + // Pass 3: claim a dict_id in the bucket, overwrite slot.second, materialize + // dict_data. The atomicAdd is shared-memory only, so no global traffic. + for (size_type slot_idx = t; slot_idx < dict_map_size; slot_idx += block_size) { + auto const slot_key = chunk_slots[slot_idx].first; + if (slot_key != KEY_SENTINEL) { + auto const frag_local = chunk_slots[slot_idx].second - f_start; + auto const loc = atomicAdd(&fragment_cursor[frag_local], 1); cudf_assert(loc < MAX_DICT_SIZE && "Number of filled slots exceeds max dict size"); - chunk.dict_data[loc] = key; - // If sorting dict page ever becomes a hard requirement, enable the following statement - // and add a dict sorting step before storing into the slot's second field. - // chunk.dict_data_idx[loc] = idx; - slot->second = loc; + chunk.dict_data[loc] = slot_key; + chunk_slots[slot_idx].second = loc; } } } @@ -299,11 +422,12 @@ void populate_chunk_hash_maps(device_span const map_storage, void collect_map_entries(device_span const map_storage, device_span chunks, + cudf::detail::device_2dspan frags, rmm::cuda_stream_view stream) { constexpr int block_size = 1024; collect_map_entries_kernel - <<>>(map_storage, chunks); + <<>>(map_storage, chunks, frags); } void get_dictionary_indices(device_span const map_storage, diff --git a/cpp/src/io/parquet/parquet_gpu.cuh b/cpp/src/io/parquet/parquet_gpu.cuh index ed79cb1ff06d..15323e8f30f4 100644 --- a/cpp/src/io/parquet/parquet_gpu.cuh +++ b/cpp/src/io/parquet/parquet_gpu.cuh @@ -97,12 +97,23 @@ void populate_chunk_hash_maps(device_span const map_storage, /** * @brief Compact dictionary hash map entries into chunk.dict_data * + * Each chunk's `dict_id`s are assigned monotonically in fragment-first- + * appearance order: values whose slot was first won by an earlier fragment + * get strictly lower `dict_id`s than values first won by a later fragment. + * This is the ordering prerequisite that lets per-page RLE bit widths shrink + * for pages that only touch values in early fragments. + * * @param map_storage Bulk hashmap storage * @param chunks Flat span of chunks to compact hash maps for + * @param frags 2D span of per-column page fragments. Used by the collect + * kernel to determine each chunk's column-relative fragment + * range; the span itself matches the one passed to + * `populate_chunk_hash_maps`. * @param stream CUDA stream to use */ void collect_map_entries(device_span const map_storage, device_span chunks, + cudf::detail::device_2dspan frags, rmm::cuda_stream_view stream); /** diff --git a/cpp/src/io/parquet/writer_impl.cu b/cpp/src/io/parquet/writer_impl.cu index 571d36133117..da633a9e4393 100644 --- a/cpp/src/io/parquet/writer_impl.cu +++ b/cpp/src/io/parquet/writer_impl.cu @@ -1409,7 +1409,7 @@ build_chunk_dictionaries(hostdevice_2dvector& chunks, chunk.dict_index = inserted_dict_index.data(); } chunks.host_to_device_async(stream); - collect_map_entries(map_storage_data, chunks.device_view().flat_view(), stream); + collect_map_entries(map_storage_data, chunks.device_view().flat_view(), frags, stream); get_dictionary_indices(map_storage_data, frags, stream); return std::pair(std::move(dict_data), std::move(dict_index)); From 5cc906ec30d4c72a61c669839b3e20330633ff77 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Thu, 23 Apr 2026 00:24:15 +0000 Subject: [PATCH 03/28] OPT: per-page variable bit-width RLE for parquet dict indices (#13995) Each data page now RLE-encodes dictionary indices using ceil(log2(page_max_dict_index + 1)) bits instead of the chunk-wide maximum. Combined with the first-appearance dict_id ordering from the prior commit, this closes the ~30% file-size gap vs Spark on moderate-cardinality INT64 and STRING workloads. Page size estimation continues to use the chunk-wide bits as a conservative upper bound; the dictionary page itself is unaffected. Made-with: Cursor --- .../io/parquet/parquet_writer_dict.cpp | 110 ++++++- cpp/src/io/parquet/chunk_dict.cu | 109 +++++- cpp/src/io/parquet/page_enc.cu | 16 +- cpp/src/io/parquet/parquet_gpu.cuh | 17 + cpp/src/io/parquet/parquet_gpu.hpp | 6 + cpp/src/io/parquet/writer_impl.cu | 8 + cpp/tests/io/parquet_misc_test.cpp | 19 +- cpp/tests/io/parquet_writer_test.cpp | 310 ++++++++++++++++++ 8 files changed, 574 insertions(+), 21 deletions(-) diff --git a/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp b/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp index 0e9214deb5ee..dcabd60b80a3 100644 --- a/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp +++ b/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp @@ -9,18 +9,26 @@ #include #include +#include +#include #include +#include +#include #include #include #include #include #include +#include + #include +#include #include #include #include +#include #include #include #include @@ -110,6 +118,51 @@ std::vector build_numeric_column(cudf::size_type num_rows, cudf::size_type pa return std::make_unique(std::move(cols)); } +// Walk the written file's offset indexes and return the per-data-page RLE +// bit-width byte for dict-encoded pages. Assumes a flat column with no +// rep/def levels (the benchmark builds INT64 no-nulls columns); under that +// layout the first byte of a V1 data page payload is the RLE dict_bits +// (PROBLEM.md §2). For V2 data pages we'd need to skip past the declared +// def/rep byte lengths before reading the bit-width, but this benchmark +// emits V1 headers so the simpler path is sufficient. +// +// The caller is expected to have already populated `chunk.offset_index` on +// the `FileMetaData` (e.g. via `hybrid_scan_reader::setup_page_index`). +// `read_parquet_footers` alone is insufficient: it only materializes the +// OffsetIndex when the file has BYTE_ARRAY columns (see `metadata::metadata` +// in reader_impl_helpers.cpp), which is not the case for this INT64 bench. +[[nodiscard]] std::vector extract_page_dict_bits( + cudf::io::datasource& source, cudf::io::parquet::FileMetaData const& footer) +{ + using namespace cudf::io::parquet; + std::vector bits; + for (auto const& rg : footer.row_groups) { + for (auto const& chunk : rg.columns) { + if (not chunk.offset_index.has_value()) { continue; } + for (auto const& pl : chunk.offset_index->page_locations) { + if (pl.offset <= 0 || pl.compressed_page_size <= 0) { continue; } + auto const buf = source.host_read(pl.offset, pl.compressed_page_size); + detail::CompactProtocolReader cp(buf->data(), buf->size()); + PageHeader hdr; + cp.read(&hdr); + auto const is_dict_encoded = + (hdr.type == PageType::DATA_PAGE && + (hdr.data_page_header.encoding == Encoding::PLAIN_DICTIONARY || + hdr.data_page_header.encoding == Encoding::RLE_DICTIONARY)) || + (hdr.type == PageType::DATA_PAGE_V2 && + (hdr.data_page_header_v2.encoding == Encoding::PLAIN_DICTIONARY || + hdr.data_page_header_v2.encoding == Encoding::RLE_DICTIONARY)); + if (not is_dict_encoded) { continue; } + // `cp` is positioned at the first byte of the page payload after the + // header thrift; that byte is the RLE bit width for dict-indexed + // pages (valid only with no rep/def levels, which holds here). + bits.push_back(cp.getb()); + } + } + } + return bits; +} + } // namespace void BM_parq_write_dict_encoding(nvbench::state& state) @@ -155,16 +208,53 @@ void BM_parq_write_dict_encoding(nvbench::state& state) mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage"); state.add_buffer_size(encoded_file_size, "encoded_file_size", "encoded_file_size"); - // TODO(phase 2): once per-page `dict_rle_bits` is plumbed into the writer - // path and exposed at the reader, replace this whole-file proxy with an - // exact min/max/mean across emitted pages. Until then we report aggregate - // bits-per-value, which tracks overall compression ratio but not the - // per-page distribution we ultimately want to shrink. - auto const total_values = static_cast(view.num_rows()) * num_cols; - if (total_values > 0) { - auto const bits_per_value_avg = - static_cast(encoded_file_size * 8) / static_cast(total_values); - state.add_element_count(bits_per_value_avg, "bits_per_value_avg"); + // Emit the per-page dictionary RLE bit-width distribution. We do an extra + // untimed write with `STATISTICS_COLUMN` so that the OffsetIndex (which + // gives us each page's byte range) is actually written; the timed write + // above intentionally uses the default stats granularity to keep the + // benchmark's encoding path unperturbed. For this benchmark workload + // (flat INT64, no nulls, V1 headers) the first byte of each data page's + // payload is the RLE bit-width used to encode dict indices. + cuio_source_sink_pair inspect_sink(io_type::FILEPATH); + { + auto const inspect_opts = + cudf::io::parquet_writer_options::builder(inspect_sink.make_sink_info(), view) + .compression(cudf::io::compression_type::NONE) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .row_group_size_rows(num_rows) + .max_page_size_rows(page_size_rows) + .max_page_size_bytes(std::size_t{64} << 20) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .build(); + cudf::io::write_parquet(inspect_opts); + } + + // Use the hybrid scan reader to build a `FileMetaData` with a fully + // materialized OffsetIndex for all column chunks (works regardless of + // column type, unlike `read_parquet_footers` which skips the page index + // when no string columns are present). + auto inspect_sources = cudf::io::make_datasources(inspect_sink.make_source_info()); + auto& inspect_ds = *inspect_sources.front(); + auto const footer_buf = cudf::io::parquet::fetch_footer_to_host(inspect_ds); + cudf::io::parquet::experimental::hybrid_scan_reader reader(*footer_buf, + cudf::io::parquet_reader_options{}); + auto const page_index_range = reader.page_index_byte_range(); + if (not page_index_range.is_empty()) { + auto const pi_buf = + cudf::io::parquet::fetch_page_index_to_host(inspect_ds, page_index_range); + reader.setup_page_index(*pi_buf); + } + auto const footer = reader.parquet_metadata(); + auto const page_bits = extract_page_dict_bits(inspect_ds, footer); + + if (not page_bits.empty()) { + auto const [min_it, max_it] = std::minmax_element(page_bits.begin(), page_bits.end()); + auto const sum = + std::accumulate(page_bits.begin(), page_bits.end(), std::uint64_t{0}); + auto const mean = static_cast(sum) / static_cast(page_bits.size()); + state.add_element_count(static_cast(*min_it), "dict_rle_bits_min"); + state.add_element_count(static_cast(*max_it), "dict_rle_bits_max"); + state.add_element_count(mean, "dict_rle_bits_mean"); } } diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index 55f3512e2f94..33c65a521b3f 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -8,13 +8,17 @@ #include #include #include +#include #include #include #include +#include #include #include +#include +#include namespace cudf::io::parquet::detail { @@ -128,8 +132,8 @@ struct map_insert_fn { // `cuco::op::min`, switch to that for exact first-fragment semantics // (PROBLEM.md §5.2 Option II). auto const fragment_hint = static_cast(blockIdx.x); - is_unique = map_insert_ref.insert( - slot_type{static_cast(val_idx), fragment_hint}); + is_unique = + map_insert_ref.insert(slot_type{static_cast(val_idx), fragment_hint}); uniq_elem_size = [&]() -> size_type { if (not is_unique) { return 0; } switch (col->physical_type) { @@ -306,8 +310,8 @@ CUDF_KERNEL void __launch_bounds__(block_size) __shared__ size_type f_start; __shared__ size_type n_frags; if (t == 0) { - f_start = static_cast(chunk.fragments - col_frag.data()); - size_type n = 0; + f_start = static_cast(chunk.fragments - col_frag.data()); + size_type n = 0; auto const total_col_frags = static_cast(col_frag.size()); while (f_start + n < total_col_frags && col_frag[f_start + n].chunk == &chunk) { ++n; @@ -373,8 +377,8 @@ CUDF_KERNEL void __launch_bounds__(block_size) auto const frag_local = chunk_slots[slot_idx].second - f_start; auto const loc = atomicAdd(&fragment_cursor[frag_local], 1); cudf_assert(loc < MAX_DICT_SIZE && "Number of filled slots exceeds max dict size"); - chunk.dict_data[loc] = slot_key; - chunk_slots[slot_idx].second = loc; + chunk.dict_data[loc] = slot_key; + chunk_slots[slot_idx].second = loc; } } } @@ -438,4 +442,97 @@ void get_dictionary_indices(device_span const map_storage, get_dictionary_indices_kernel <<>>(map_storage, frags); } + +namespace { + +// Warps per block for `compute_page_dict_rle_bits_kernel`. Sized so that each +// block pulls several pages through the reduce, which amortizes the block-level +// overhead over 4 independent warp-level reductions without oversubscribing +// shared memory or hurting occupancy. +constexpr int kDictRleBitsWarpsPerBlock = 4; +constexpr int kDictRleBitsBlockSize = kDictRleBitsWarpsPerBlock * cudf::detail::warp_size; + +// One warp per data page. Each warp strides through its page's slice of +// `chunk->dict_index`, computes the max over *valid* rows only, and stores +// `max(NumRequiredBits(page_max), 1)` into `page.dict_rle_bits`. +// +// Why warp-per-page instead of a `cub::DeviceSegmentedReduce::Max` call per +// chunk (PHASE_2_VARIABLE_BITS.md §2.2.3): segments are small (typically 1k- +// 100k elements per page), so one CUB call per chunk would be launch-overhead +// bound across workloads with many chunks. A single kernel with one warp per +// page keeps the reduction entirely in registers + shuffle and issues exactly +// one launch for the whole writer. +// +// Pages that are skipped (dictionary page itself, non-dict chunks, BOOLEAN +// columns whose dict_bits is always 1) retain the `chunk->dict_rle_bits` +// value that `gpuInitPages` wrote, so the encoder falls back to chunk-wide +// behavior for them. +CUDF_KERNEL void __launch_bounds__(kDictRleBitsBlockSize) + compute_page_dict_rle_bits_kernel(device_span pages) +{ + constexpr auto warp_size = cudf::detail::warp_size; + auto const warp_lane = static_cast(threadIdx.x % warp_size); + auto const warp_id = static_cast(threadIdx.x / warp_size); + auto const page_idx = static_cast(blockIdx.x) * kDictRleBitsWarpsPerBlock + warp_id; + + __shared__ + typename cub::WarpReduce::TempStorage reduce_storage[kDictRleBitsWarpsPerBlock]; + + if (page_idx >= static_cast(pages.size())) { return; } + + auto& page = pages[page_idx]; + auto const* chunk = page.chunk; + // Non-dict chunk: `dict_rle_bits` is unused by the encoder; skip. + if (not chunk->use_dictionary) { return; } + // Dictionary page itself does not encode dict_indices; skip. + if (page.page_type == PageType::DICTIONARY_PAGE) { return; } + + auto const* col = chunk->col_desc; + // BOOLEAN columns emit dict_bits=1 through a separate code path in + // `gpuEncodeDictPages`, independent of `dict_rle_bits`. Leave the field at + // its chunk-wide init value (it is ignored for booleans). + if (col->physical_type == Type::BOOLEAN) { return; } + + auto const chunk_start_val = row_to_value_idx(chunk->start_row, *col); + auto const page_start_val = row_to_value_idx(page.start_row, *col); + auto const page_num_leaf_values = static_cast(page.num_leaf_values); + auto const begin = page_start_val - chunk_start_val; + auto const end = begin + page_num_leaf_values; + + auto const* dict_index = chunk->dict_index; + column_device_view const& leaf_col = *col->leaf_column; + auto const leaf_size = leaf_col.size(); + + // Accumulate per-lane max. Null rows leave `dict_index` undefined; gate + // the read with the column's validity bitmap to avoid pulling garbage + // bits into the max. + size_type lane_max = 0; + for (size_type i = begin + warp_lane; i < end; i += warp_size) { + auto const val_idx = chunk_start_val + i; + if (val_idx < leaf_size && leaf_col.is_valid(val_idx)) { + lane_max = max(lane_max, dict_index[i]); + } + } + + auto const page_max = + cub::WarpReduce(reduce_storage[warp_id]).Reduce(lane_max, cuda::maximum{}); + + if (warp_lane == 0) { + // Floor at 1 to match the chunk-wide convention (all-null pages still + // emit a 1-bit RLE preamble; see `writer_impl.cu`'s `std::max(..., 1)`). + auto const nbits = max(cuda::std::bit_width(static_cast(page_max)), 1); + page.dict_rle_bits = static_cast(nbits); + } +} + +} // namespace + +void compute_per_page_dict_rle_bits(device_span pages, rmm::cuda_stream_view stream) +{ + if (pages.empty()) { return; } + auto const num_blocks = cudf::util::div_rounding_up_safe(static_cast(pages.size()), + kDictRleBitsWarpsPerBlock); + compute_page_dict_rle_bits_kernel<<>>( + pages); +} } // namespace cudf::io::parquet::detail diff --git a/cpp/src/io/parquet/page_enc.cu b/cpp/src/io/parquet/page_enc.cu index fcd90f6f710e..909cb60f2eee 100644 --- a/cpp/src/io/parquet/page_enc.cu +++ b/cpp/src/io/parquet/page_enc.cu @@ -677,6 +677,9 @@ CUDF_KERNEL void __launch_bounds__(128) page_g.num_rows = ck_g.num_dict_entries; page_g.num_leaf_values = ck_g.num_dict_entries; page_g.num_values = ck_g.num_dict_entries; // TODO: shouldn't matter for dict page + // Unused for the dictionary page itself (guarded in gpuEncodeDictPages), but + // initialized for consistency so the field is never observed zero on a populated page. + page_g.dict_rle_bits = ck_g.dict_rle_bits; page_offset += util::round_up_unsafe(page_g.max_hdr_size + page_g.max_data_size, page_align); if (not comp_page_sizes.empty()) { @@ -778,6 +781,10 @@ CUDF_KERNEL void __launch_bounds__(128) page_g.data_size = 0; page_g.comp_data_size = 0; page_g.is_compressed = false; + // Conservative initial value: the chunk-wide bit width. A subsequent pass + // (`compute_per_page_dict_rle_bits`) tightens this to the page-local max + // once page boundaries are fixed and dict_index has been materialized. + page_g.dict_rle_bits = ck_g.dict_rle_bits; page_g.max_hdr_size = max_data_page_hdr_size; // Max size excluding statistics if (ck_g.stats) { uint32_t stats_hdr_len = 16; @@ -1909,9 +1916,16 @@ CUDF_KERNEL void __launch_bounds__(block_size, 8) }(); // TODO assert dict_bits >= 0 + // + // For data pages we use `page.dict_rle_bits`, which `compute_per_page_dict_rle_bits` + // tightens to `ceil(log2(page_max_dict_index + 1))`, not the chunk-wide upper bound + // stored in `chunk.dict_rle_bits`. This is the encode-side half of the per-page + // variable-bit-width optimization (PHASE_2_VARIABLE_BITS.md). Page-size estimation + // in `gpuInitPages` intentionally keeps using `chunk.dict_rle_bits` so buffer sizing + // remains a conservative upper bound -- PROBLEM.md §6. auto const dict_bits = (physical_type == Type::BOOLEAN) ? 1 : (s->ck.use_dictionary and s->page.page_type != PageType::DICTIONARY_PAGE) - ? s->ck.dict_rle_bits + ? s->page.dict_rle_bits : -1; if (t == 0) { uint8_t* dst = s->cur; diff --git a/cpp/src/io/parquet/parquet_gpu.cuh b/cpp/src/io/parquet/parquet_gpu.cuh index 15323e8f30f4..cbd1b4cf91a9 100644 --- a/cpp/src/io/parquet/parquet_gpu.cuh +++ b/cpp/src/io/parquet/parquet_gpu.cuh @@ -133,4 +133,21 @@ void get_dictionary_indices(device_span const map_storage, cudf::detail::device_2dspan frags, rmm::cuda_stream_view stream); +/** + * @brief Tighten each data page's `dict_rle_bits` to the minimum width required + * by the max `dict_index` observed in that page's rows. + * + * Must be invoked after `InitEncoderPages` has finalized page boundaries *and* + * `get_dictionary_indices` has materialized per-chunk `dict_index` arrays, but + * before `EncodePages` reads `page.dict_rle_bits`. Pages that do not use + * dictionary encoding (non-dict chunks, the dictionary page itself, BOOLEAN + * columns) are skipped and keep the `chunk->dict_rle_bits` fallback that + * `gpuInitPages` wrote during page initialization. + * + * @param pages Device span of encoder pages. Field `dict_rle_bits` is written. + * @param stream CUDA stream to use + */ +void compute_per_page_dict_rle_bits(device_span pages, + rmm::cuda_stream_view stream); + } // namespace cudf::io::parquet::detail diff --git a/cpp/src/io/parquet/parquet_gpu.hpp b/cpp/src/io/parquet/parquet_gpu.hpp index abea80f9d764..7397a063ba9d 100644 --- a/cpp/src/io/parquet/parquet_gpu.hpp +++ b/cpp/src/io/parquet/parquet_gpu.hpp @@ -649,6 +649,12 @@ struct EncPage { Encoding encoding; //!< Encoding used for page data uint16_t num_fragments; //!< Number of fragments in page bool is_compressed; //!< Whether this page is compressed (for V2 page-level compression) + uint8_t dict_rle_bits; //!< RLE bit width for this data page's dict indices. + //!< Initialized to `chunk->dict_rle_bits` (the conservative + //!< chunk-wide bound) and then tightened by + //!< `compute_per_page_dict_rle_bits` to + //!< `max(NumRequiredBits(page_max_dict_index), 1)`. Unused + //!< for the dictionary page itself and for non-dict pages. [[nodiscard]] CUDF_HOST_DEVICE constexpr bool is_v2() const { diff --git a/cpp/src/io/parquet/writer_impl.cu b/cpp/src/io/parquet/writer_impl.cu index da633a9e4393..e9e95a0828cf 100644 --- a/cpp/src/io/parquet/writer_impl.cu +++ b/cpp/src/io/parquet/writer_impl.cu @@ -2118,6 +2118,14 @@ auto convert_table_to_parquet_data(table_input_metadata& table_meta, max_page_size_rows, write_v2_headers, stream); + + // Now that page boundaries are finalized and `chunk->dict_index` has been + // materialized by `build_chunk_dictionaries`, tighten each data page's + // `dict_rle_bits` from the chunk-wide conservative bound to the page-local + // minimum required width. Page-size estimation intentionally keeps the + // chunk-wide value (see PROBLEM.md §6): estimation is upstream of this + // pass and must remain conservative so page buffers never overflow. + compute_per_page_dict_rle_bits({pages.data(), pages.size()}, stream); } // Check device write support for all chunks and initialize bounce_buffer. diff --git a/cpp/tests/io/parquet_misc_test.cpp b/cpp/tests/io/parquet_misc_test.cpp index 4bcf31146a33..0de0628ed0a1 100644 --- a/cpp/tests/io/parquet_misc_test.cpp +++ b/cpp/tests/io/parquet_misc_test.cpp @@ -169,10 +169,21 @@ TEST_P(ParquetSizedTest, DictionaryTest) }); EXPECT_TRUE(used_dict); - // and check that the correct number of bits was used - auto const oi = read_offset_index(source, fmd.row_groups.front().columns.front()); - auto const nbits = read_dict_bits(source, oi.page_locations.front()); - EXPECT_EQ(nbits, GetParam()); + // and check that the correct number of bits was used. Phase 2 variable-bit-width + // encoding lets each page emit `ceil(log2(page_max_dict_index + 1))` bits rather + // than the chunk-wide maximum, so individual pages may use *fewer* bits than + // `GetParam()`. The chunk-wide maximum must still be reached by at least one page + // (the one that references the last-assigned dict_id), so we verify: + // - every page's bit width is <= GetParam() (upper bound is respected), and + // - max page bit width == GetParam() (the chunk-wide bound is tight). + auto const oi = read_offset_index(source, fmd.row_groups.front().columns.front()); + int max_nbits = 0; + for (auto const& pl : oi.page_locations) { + auto const nbits = read_dict_bits(source, pl); + EXPECT_LE(nbits, GetParam()); + max_nbits = std::max(max_nbits, nbits); + } + EXPECT_EQ(max_nbits, GetParam()); } /////////////////////// diff --git a/cpp/tests/io/parquet_writer_test.cpp b/cpp/tests/io/parquet_writer_test.cpp index 5e84299a1541..15a40286d828 100644 --- a/cpp/tests/io/parquet_writer_test.cpp +++ b/cpp/tests/io/parquet_writer_test.cpp @@ -22,13 +22,17 @@ #include #include +#include + #include #include #include +#include #include #include +#include using cudf::test::iterators::no_nulls; @@ -1106,6 +1110,312 @@ TEST_F(ParquetWriterTest, SingleValueDictionaryTest) EXPECT_EQ(nbits, expected_bits); } +// Phase 2 per-page variable-bit-width RLE coverage (PHASE_2_VARIABLE_BITS.md §2.4). +// The four tests below exercise the new code path from three angles: +// (1) decoder correctness across cardinalities and physical types +// (integers / strings / list), +// (2) the chunk-wide upper bound is respected on every page, +// (3) the optimization actually fires -- file size drops below the pre-Phase-2 +// chunk-wide number on the benchmark workload. + +// `ceil(log2(max_dict_index + 1))`, floored at 1. Matches what +// `compute_page_dict_rle_bits_kernel` writes device-side (`32 - +// countl_zero(max)`, floored at 1) and the chunk-wide convention in +// `build_chunk_dictionaries` (all-null pages still emit a 1-bit RLE preamble, +// see writer_impl.cu's `std::max(..., 1)`). +namespace { +[[nodiscard]] int num_required_bits(uint32_t v) { return std::max(std::bit_width(v), 1); } + +// Returns true iff the column chunk was actually written with dictionary +// encoding. The writer may silently fall back to PLAIN for sparse-cardinality +// inputs even under `dictionary_policy::ALWAYS` (e.g., when the dict page +// would dwarf the data page), so tests that inspect per-page bit widths must +// check this before dereferencing `read_dict_bits` output. +[[nodiscard]] bool chunk_used_dictionary(cudf::io::parquet::ColumnChunk const& chunk) +{ + for (auto const enc : chunk.meta_data.encodings) { + if (enc == cudf::io::parquet::Encoding::PLAIN_DICTIONARY or + enc == cudf::io::parquet::Encoding::RLE_DICTIONARY) { + return true; + } + } + return false; +} +} // namespace + +// Round-trip INT64 dictionary-encoded columns at three cardinalities. The +// cardinalities are chosen to straddle the 10-bit / 16-bit / 20-bit RLE widths +// so the reader has to correctly consume per-page preambles of different +// widths. Chunk-wide-width encoding would pass round-trip as well, so this is +// primarily a "no reader mismatch under variable widths" gate; the upper-bound +// assertion catches an accidental regression that emits *more* bits per page +// than the chunk actually needs. +TEST_F(ParquetWriterTest, VariableBitWidthRoundTripIntegers) +{ + constexpr cudf::size_type nrows = 200'000; + + for (auto const cardinality : {10'000, 64'000, 1'000'000}) { + std::mt19937 rng{0xFEEDFACE}; + std::uniform_int_distribution dist(0, cardinality - 1); + std::vector values(nrows); + std::generate(values.begin(), values.end(), [&] { return dist(rng); }); + + auto const col = cudf::test::fixed_width_column_wrapper(values.begin(), values.end()); + auto const expected = table_view{{col}}; + + auto buffer = std::vector{}; + cudf::io::parquet_writer_options out_opts = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&buffer}, expected) + .compression(cudf::io::compression_type::NONE) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .row_group_size_rows(nrows) + .max_page_size_rows(nrows / 4); + cudf::io::write_parquet(out_opts); + + auto const buffer_span = + cudf::host_span(reinterpret_cast(buffer.data()), buffer.size()); + cudf::io::parquet_reader_options in_opts = + cudf::io::parquet_reader_options::builder(cudf::io::source_info(buffer_span)); + auto const result = cudf::io::read_parquet(in_opts); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); + + auto const source = cudf::io::datasource::create(cudf::host_span{buffer_span}); + cudf::io::parquet::FileMetaData fmd; + read_footer(source, &fmd); + + auto const& chunk = fmd.row_groups.front().columns.front(); + // Under `dictionary_policy::ALWAYS` the writer still falls back to PLAIN + // when the dict page would exceed the row data (very sparse cardinality + // relative to row count). Treat that as out-of-scope for bit-width + // checking but still validate round-trip correctness above. + if (not chunk_used_dictionary(chunk)) { continue; } + + auto const chunk_wide_max_bits = num_required_bits(cardinality - 1); + auto const oi = read_offset_index(source, chunk); + ASSERT_GT(oi.page_locations.size(), 1u) << "cardinality=" << cardinality; + for (auto const& pl : oi.page_locations) { + auto const nbits = read_dict_bits(source, pl); + EXPECT_GE(nbits, 1) << "cardinality=" << cardinality; + EXPECT_LE(nbits, chunk_wide_max_bits) << "cardinality=" << cardinality; + } + } +} + +// Round-trip string dictionary encoding at moderate cardinality. Strings take a +// different write path through `build_chunk_dictionaries` (variable-length hash +// keys, dict page is an array of `string_index_pair`), so we cover the encode +// path with a separate fixture rather than folding it into the integer test. +TEST_F(ParquetWriterTest, VariableBitWidthRoundTripStrings) +{ + constexpr cudf::size_type nrows = 100'000; + constexpr cudf::size_type cardinality = 5'000; + + std::mt19937 rng{0xDEADBEEF}; + std::uniform_int_distribution dist(0, cardinality - 1); + std::vector values(nrows); + std::generate(values.begin(), values.end(), [&] { return "str_" + std::to_string(dist(rng)); }); + + auto const col = cudf::test::strings_column_wrapper(values.begin(), values.end()); + auto const expected = table_view{{col}}; + + auto buffer = std::vector{}; + cudf::io::parquet_writer_options out_opts = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&buffer}, expected) + .compression(cudf::io::compression_type::NONE) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .row_group_size_rows(nrows) + .max_page_size_rows(nrows / 4); + cudf::io::write_parquet(out_opts); + + auto const buffer_span = + cudf::host_span(reinterpret_cast(buffer.data()), buffer.size()); + cudf::io::parquet_reader_options in_opts = + cudf::io::parquet_reader_options::builder(cudf::io::source_info(buffer_span)); + auto const result = cudf::io::read_parquet(in_opts); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); + + auto const source = cudf::io::datasource::create(cudf::host_span{buffer_span}); + cudf::io::parquet::FileMetaData fmd; + read_footer(source, &fmd); + + auto const chunk_wide_max_bits = num_required_bits(cardinality - 1); + auto const oi = read_offset_index(source, fmd.row_groups.front().columns.front()); + ASSERT_GT(oi.page_locations.size(), 1u); + for (auto const& pl : oi.page_locations) { + auto const nbits = read_dict_bits(source, pl); + EXPECT_GE(nbits, 1); + EXPECT_LE(nbits, chunk_wide_max_bits); + } +} + +// Round-trip list dictionary encoding. Nested types exercise the +// `row_to_value_idx` path inside `compute_page_dict_rle_bits_kernel`: the +// per-page max must be taken over *leaf* values, not rows, and the kernel must +// correctly translate `page.start_row`/`page.num_leaf_values` into leaf-space +// offsets. A mis-translation here would produce a bit width too small to cover +// some leaf dict_index, and the reader would mis-decode the page -- so the +// round-trip comparison below is the real gate. We deliberately don't inspect +// `read_dict_bits` here because the helper assumes a flat column layout (dict +// RLE stream at byte 0 of the page payload); for list the page payload +// starts with rep+def levels and dereferencing byte 0 as a bit width is +// meaningless. +TEST_F(ParquetWriterTest, VariableBitWidthRoundTripLists) +{ + constexpr cudf::size_type num_lists = 20'000; + constexpr cudf::size_type cardinality = 1'024; + + std::mt19937 rng{0xCAFEF00D}; + std::uniform_int_distribution list_len_dist(0, 5); + std::uniform_int_distribution val_dist(0, cardinality - 1); + + std::vector leaf_values; + std::vector offsets{0}; + leaf_values.reserve(num_lists * 3); + offsets.reserve(num_lists + 1); + for (cudf::size_type i = 0; i < num_lists; ++i) { + auto const list_len = list_len_dist(rng); + for (int j = 0; j < list_len; ++j) { + leaf_values.push_back(val_dist(rng)); + } + offsets.push_back(static_cast(leaf_values.size())); + } + + auto leaf_col = + cudf::test::fixed_width_column_wrapper(leaf_values.begin(), leaf_values.end()) + .release(); + auto offsets_col = + cudf::test::fixed_width_column_wrapper(offsets.begin(), offsets.end()) + .release(); + auto list_col = cudf::make_lists_column( + num_lists, std::move(offsets_col), std::move(leaf_col), 0, rmm::device_buffer{}); + auto const expected = table_view{{*list_col}}; + + auto buffer = std::vector{}; + cudf::io::parquet_writer_options out_opts = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&buffer}, expected) + .compression(cudf::io::compression_type::NONE) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .row_group_size_rows(num_lists) + .max_page_size_rows(num_lists / 4); + cudf::io::write_parquet(out_opts); + + auto const buffer_span = + cudf::host_span(reinterpret_cast(buffer.data()), buffer.size()); + cudf::io::parquet_reader_options in_opts = + cudf::io::parquet_reader_options::builder(cudf::io::source_info(buffer_span)); + auto const result = cudf::io::read_parquet(in_opts); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); + + auto const source = cudf::io::datasource::create(cudf::host_span{buffer_span}); + cudf::io::parquet::FileMetaData fmd; + read_footer(source, &fmd); + EXPECT_TRUE(chunk_used_dictionary(fmd.row_groups.front().columns.front())); +} + +// End-to-end file-size gate. Reproduces the `parquet_write_dict_encoding` +// benchmark workload (8 "common" pages touching only 64 frequent values + 2 +// "rare" pages touching values 64..63999) and asserts: +// (a) common pages bit-pack at their page-local minimum width (`nbits <= +// num_required_bits(frequent_set_size - 1)`), +// (b) at least one page reaches the chunk-wide width (so dict_rle_bits is +// still the upper bound, not silently clobbered), +// (c) the resulting file is strictly smaller than the Phase 1 baseline +// (2,491,156 bytes at num_rows = 1,000,000). We use the benchmark's exact +// shape but a smaller row count (200K rows, 20K rows/page) so the test +// finishes in well under a second; the proportional savings are the same. +// See PHASE_2_VARIABLE_BITS.md §2.4 for the design rationale. +TEST_F(ParquetWriterTest, VariableBitWidthSmallerThanChunkWide) +{ + constexpr cudf::size_type pages_per_chunk = 10; + constexpr cudf::size_type hot_pages = pages_per_chunk - 2; + constexpr cudf::size_type page_size_rows = 20'000; + constexpr cudf::size_type num_rows = pages_per_chunk * page_size_rows; + constexpr cudf::size_type cardinality = 64'000; + constexpr cudf::size_type frequent_set_size = 64; + + std::mt19937 rng{0xC0DEFACE}; + std::uniform_int_distribution freq_dist(0, frequent_set_size - 1); + std::uniform_int_distribution rare_dist(frequent_set_size, cardinality - 1); + + std::vector values(num_rows); + cudf::size_type const threshold = hot_pages * page_size_rows; + for (cudf::size_type i = 0; i < num_rows; ++i) { + values[i] = i < threshold ? freq_dist(rng) : rare_dist(rng); + } + + auto const col = cudf::test::fixed_width_column_wrapper(values.begin(), values.end()); + auto const expected = table_view{{col}}; + + auto buffer = std::vector{}; + cudf::io::parquet_writer_options out_opts = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&buffer}, expected) + .compression(cudf::io::compression_type::NONE) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .row_group_size_rows(num_rows) + .max_page_size_rows(page_size_rows) + .max_page_size_bytes(std::size_t{64} << 20); + cudf::io::write_parquet(out_opts); + + auto const buffer_span = + cudf::host_span(reinterpret_cast(buffer.data()), buffer.size()); + cudf::io::parquet_reader_options in_opts = + cudf::io::parquet_reader_options::builder(cudf::io::source_info(buffer_span)); + auto const result = cudf::io::read_parquet(in_opts); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); + + auto const source = cudf::io::datasource::create(cudf::host_span{buffer_span}); + cudf::io::parquet::FileMetaData fmd; + read_footer(source, &fmd); + + ASSERT_TRUE(chunk_used_dictionary(fmd.row_groups.front().columns.front())); + + // Chunk-wide upper bound computed from the full cardinality. The writer's + // `dict_rle_bits` is `NumRequiredBits(num_dict_entries - 1)`, and under + // uniform sampling of `[0, cardinality)` the actual number of distinct + // values drawn may land a bit below cardinality -- so we use this value as + // a *ceiling* rather than an equality target. + auto const chunk_wide_max_bits = num_required_bits(cardinality - 1); + auto const frequent_max_bits = num_required_bits(frequent_set_size - 1); // 6 + + auto const oi = read_offset_index(source, fmd.row_groups.front().columns.front()); + ASSERT_EQ(oi.page_locations.size(), static_cast(pages_per_chunk)); + + int common_page_count = 0; + int rare_page_count = 0; + int max_observed_bits = 0; + for (auto const& pl : oi.page_locations) { + auto const nbits = read_dict_bits(source, pl); + EXPECT_GE(nbits, 1); + EXPECT_LE(nbits, chunk_wide_max_bits); + max_observed_bits = std::max(max_observed_bits, nbits); + if (nbits <= frequent_max_bits) { + ++common_page_count; + } else { + ++rare_page_count; + } + } + // Under Phase 1 ordering the first `hot_pages` pages reference only the + // frequent-set dict_ids `[0, 64)`, so they should bit-pack to <= 6 bits. + // The 2 rare pages reference dict_ids outside that range and must use + // strictly more bits -- this is the whole optimization in one assertion. + EXPECT_EQ(common_page_count, hot_pages); + EXPECT_EQ(rare_page_count, pages_per_chunk - hot_pages); + EXPECT_GT(max_observed_bits, frequent_max_bits); + + // Pre-Phase-2 (chunk-wide 16 bits on every page) baseline at this workload + // scales linearly from the 1M-row benchmark number (2,491,156 bytes): + // 2,491,156 * (200'000 / 1'000'000) = 498,231 bytes + // With per-page variable widths the 8 common pages drop from 16 to <=6 bits + // and should shave ~250 KB (8 pages * 20k values * (16-6) bits / 8). + // Use a conservative bound of 450,000 bytes to leave headroom for future + // encoder improvements. + EXPECT_LT(buffer.size(), 450'000u); +} + TEST_F(ParquetWriterTest, DictionaryNeverTest) { constexpr unsigned int nrows = 1'000U; From 5b0758436f392410e4cbdb4d879ae1febc345389 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Thu, 23 Apr 2026 21:34:06 +0000 Subject: [PATCH 04/28] Update benchmark --- .../io/parquet/parquet_writer_dict.cpp | 322 ++++++++---------- 1 file changed, 147 insertions(+), 175 deletions(-) diff --git a/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp b/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp index dcabd60b80a3..f68b51565bf1 100644 --- a/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp +++ b/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp @@ -3,6 +3,8 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include "io/parquet/compact_protocol_reader.hpp" + #include #include @@ -15,102 +17,90 @@ #include #include #include -#include #include #include #include -#include - #include #include -#include -#include -#include #include #include #include #include namespace { -// Dict-encoding-focused Parquet writer benchmark. Built to isolate the two -// encoder pathologies reported in https://github.com/rapidsai/cudf/issues/13995: -// -// 1. `dict_id` assignment order. cuDF's `collect_map_entries_kernel` walks -// the hash map in slot order, so the `dict_id` assigned to a value is -// uncorrelated with how early it appears in the column. Early pages -// therefore end up referencing high-indexed entries, defeating any -// per-page bit-width savings. -// 2. Chunk-wide RLE bit width. `build_chunk_dictionaries` derives one -// `dict_rle_bits = ceil(log2(num_dict_entries))` for the entire chunk, -// so a page touching only a handful of entries still bit-packs at the -// chunk-wide width. -// -// Workload shape (deliberately the simplest construction that exposes both -// problems on a single row group / single chunk): -// -// * 1 INT64 column, 1 row group = 1 chunk, 10 pages per chunk. -// INT64 is the type reported in the issue; fixed width keeps the -// per-row byte count deterministic so the file-size delta is driven -// entirely by how the encoder packs the dict-index stream. -// * cardinality = 64,000 (ceil(log2) = 16 bits chunk-wide). -// * Pages 0..`hot_pages-1` (the "common" pages) draw uniformly from a -// small "frequent set" of `frequent_set_size` values shared across -// all common pages. -// * Pages `hot_pages..pages_per_chunk-1` (the "rare" pages) draw -// uniformly from the remaining `cardinality - frequent_set_size` -// values. These never appear in any common page, so only the rare -// pages force the chunk-wide bit width up. -// -// Under an ideal encoder (first-appearance ordering + per-page bit width) -// the common pages need only `ceil(log2(frequent_set_size))` bits per -// value while the rare pages need `ceil(log2(cardinality))` bits. With -// the constants below (frequent_set_size = 64), that is 6 bits vs. 16 -// bits, saving ~10 bits/value on every common-page row. The current -// cuDF encoder produces 16 bits on every page, so the headline -// `encoded_file_size` metric cleanly resolves the optimization target. - -// Distribution shape — see file header. These are intentionally fixed so that -// the encoded-file-size number is directly comparable across phases. Adjust -// together with the phase baseline if the shape changes. -constexpr cudf::size_type num_cols = 1; -constexpr cudf::size_type cardinality = 64'000; -constexpr cudf::size_type pages_per_chunk = 10; -constexpr cudf::size_type hot_pages = pages_per_chunk - 2; -constexpr cudf::size_type frequent_set_size = 64; -constexpr std::uint32_t dict_rng_seed = 0xC0DEFACE; - -// Build the row-to-dictionary-index mapping on host for a single column. -// Index space (disjoint, exactly covers [0, cardinality)): -// * frequent set: [0, frequent_set_size) -// * rare set: [frequent_set_size, cardinality) + +constexpr auto frequent_pages_ratio = + 0.8; ///< 80% of the pages will only contain elements from the frequent set + +/** + * @brief Build a numeric column such that certain pages only contain elements from the frequent set + and others only contain elements from the rare set + * + * @tparam T Element type of the generated column values + * @param num_rows Total number of rows + * @param page_size_rows Number of rows per page + * @param cardinality Total number of distinct values + * @param frequent_set_ratio Fraction of `cardinality` assigned to the frequent set + * @return Constructed column values + */ template -std::vector build_numeric_column(cudf::size_type num_rows, cudf::size_type page_size_rows) +std::vector build_numeric_column(cudf::size_type num_rows, + cudf::size_type page_size_rows, + cudf::size_type cardinality, + double frequent_set_ratio) { - CUDF_EXPECTS(num_rows == pages_per_chunk * page_size_rows, - "num_rows must equal pages_per_chunk * page_size_rows"); - CUDF_EXPECTS(frequent_set_size < cardinality, - "cardinality must leave room for a nonempty rare set"); + static constexpr auto dict_rng_seed = 0xC0DEFACE; + + CUDF_EXPECTS(frequent_set_ratio > 0.0 and frequent_set_ratio < 1.0, + "frequent_set_ratio must be between 0.0 and 1.0"); + CUDF_EXPECTS(num_rows % page_size_rows == 0, "num_rows must be a multiple of page_size_rows"); + static_assert(frequent_pages_ratio > 0.0 and frequent_pages_ratio < 1.0, + "hot_pages_ratio must be between 0.0 and 1.0"); + + auto const total_pages = num_rows / page_size_rows; + auto const frequent_set_threshold = + static_cast(total_pages * frequent_pages_ratio) * page_size_rows; + + auto const frequent_set_size = + static_cast(static_cast(cardinality) * frequent_set_ratio); std::mt19937 rng{dict_rng_seed}; std::uniform_int_distribution freq_dist(0, frequent_set_size - 1); std::uniform_int_distribution rare_dist(frequent_set_size, cardinality - 1); - auto const threshold = hot_pages * page_size_rows; cudf::size_type row_idx = 0; std::vector values(num_rows); std::generate_n(values.begin(), num_rows, [&]() { - return row_idx++ < threshold ? static_cast(freq_dist(rng)) : static_cast(rare_dist(rng)); + return row_idx++ < frequent_set_threshold ? static_cast(freq_dist(rng)) + : static_cast(rare_dist(rng)); }); return values; } +/** + * @brief Build a table with a single INT64 column + * + * @tparam reverse_order Whether to reverse the order of the values + * @param num_rows Number of rows + * @param page_size_rows Number of rows per page + * @param cardinality Total number of distinct values + * @param frequent_set_ratio Fraction of `cardinality` assigned to the frequent set + * @return std::unique_ptr + */ +template [[nodiscard]] std::unique_ptr build_table(cudf::size_type num_rows, - cudf::size_type page_size_rows) + cudf::size_type page_size_rows, + cudf::size_type cardinality, + double frequent_set_ratio) { - auto const values = build_numeric_column(num_rows, page_size_rows); + constexpr cudf::size_type num_cols = 1; + + auto values = + build_numeric_column(num_rows, page_size_rows, cardinality, frequent_set_ratio); + if constexpr (reverse_order) { std::reverse(values.begin(), values.end()); } std::vector> cols; cols.reserve(num_cols); cols.emplace_back( @@ -118,44 +108,43 @@ std::vector build_numeric_column(cudf::size_type num_rows, cudf::size_type pa return std::make_unique(std::move(cols)); } -// Walk the written file's offset indexes and return the per-data-page RLE -// bit-width byte for dict-encoded pages. Assumes a flat column with no -// rep/def levels (the benchmark builds INT64 no-nulls columns); under that -// layout the first byte of a V1 data page payload is the RLE dict_bits -// (PROBLEM.md §2). For V2 data pages we'd need to skip past the declared -// def/rep byte lengths before reading the bit-width, but this benchmark -// emits V1 headers so the simpler path is sufficient. -// -// The caller is expected to have already populated `chunk.offset_index` on -// the `FileMetaData` (e.g. via `hybrid_scan_reader::setup_page_index`). -// `read_parquet_footers` alone is insufficient: it only materializes the -// OffsetIndex when the file has BYTE_ARRAY columns (see `metadata::metadata` -// in reader_impl_helpers.cpp), which is not the case for this INT64 bench. -[[nodiscard]] std::vector extract_page_dict_bits( - cudf::io::datasource& source, cudf::io::parquet::FileMetaData const& footer) +/** + * @brief Compute per-page RLE bit widths for dictionary-encoded pages from the parquet page index + * + * Assumption: All parquet pages are dictionary-encoded, no nulls, no rep/def levels + * + * @param source Datasource + * @param footer File metadata + * @return Vector of number of bits per page for dictionary-encoded pages + */ +[[nodiscard]] std::vector compute_page_dict_bits(cudf::io::datasource& source, + cudf::io::parquet::FileMetaData const& footer) { using namespace cudf::io::parquet; + std::vector bits; + for (auto const& rg : footer.row_groups) { for (auto const& chunk : rg.columns) { if (not chunk.offset_index.has_value()) { continue; } - for (auto const& pl : chunk.offset_index->page_locations) { - if (pl.offset <= 0 || pl.compressed_page_size <= 0) { continue; } - auto const buf = source.host_read(pl.offset, pl.compressed_page_size); - detail::CompactProtocolReader cp(buf->data(), buf->size()); - PageHeader hdr; - cp.read(&hdr); + for (auto const& page_loc : chunk.offset_index->page_locations) { + if (page_loc.offset <= 0 or page_loc.compressed_page_size <= 0) { continue; } + auto const buffer = source.host_read(page_loc.offset, page_loc.compressed_page_size); + detail::CompactProtocolReader cp(buffer->data(), buffer->size()); + PageHeader page_header; + cp.read(&page_header); + // Check if the page is dictionary-encoded. auto const is_dict_encoded = - (hdr.type == PageType::DATA_PAGE && - (hdr.data_page_header.encoding == Encoding::PLAIN_DICTIONARY || - hdr.data_page_header.encoding == Encoding::RLE_DICTIONARY)) || - (hdr.type == PageType::DATA_PAGE_V2 && - (hdr.data_page_header_v2.encoding == Encoding::PLAIN_DICTIONARY || - hdr.data_page_header_v2.encoding == Encoding::RLE_DICTIONARY)); + (page_header.type == PageType::DATA_PAGE and + (page_header.data_page_header.encoding == Encoding::PLAIN_DICTIONARY or + page_header.data_page_header.encoding == Encoding::RLE_DICTIONARY)) or + (page_header.type == PageType::DATA_PAGE_V2 and + (page_header.data_page_header_v2.encoding == Encoding::PLAIN_DICTIONARY or + page_header.data_page_header_v2.encoding == Encoding::RLE_DICTIONARY)); if (not is_dict_encoded) { continue; } // `cp` is positioned at the first byte of the page payload after the // header thrift; that byte is the RLE bit width for dict-indexed - // pages (valid only with no rep/def levels, which holds here). + // pages (valid only with no rep/def levels). bits.push_back(cp.getb()); } } @@ -167,91 +156,70 @@ std::vector build_numeric_column(cudf::size_type num_rows, cudf::size_type pa void BM_parq_write_dict_encoding(nvbench::state& state) { - auto const num_rows = static_cast(state.get_int64("num_rows")); - auto const page_size_rows = num_rows / pages_per_chunk; - - auto const tbl = build_table(num_rows, page_size_rows); - auto const view = tbl->view(); - - std::size_t encoded_file_size = 0; + auto const num_rows = static_cast(state.get_int64("num_rows")); + auto const reverse_order = static_cast(state.get_int64("reverse_order")); + auto const cardinality = static_cast(state.get_int64("cardinality")); + auto const frequent_set_ratio = static_cast(state.get_float64("frequent_set_ratio")); + auto const page_size_rows = static_cast(state.get_int64("page_size_rows")); + + CUDF_EXPECTS(page_size_rows <= num_rows and num_rows % page_size_rows == 0, + "num_rows must be a multiple of page_size_rows"); + + cuio_source_sink_pair source_sink(io_type::FILEPATH); + + auto const table = [&]() { + if (reverse_order) { + return build_table(num_rows, page_size_rows, cardinality, frequent_set_ratio); + } else { + return build_table(num_rows, page_size_rows, cardinality, frequent_set_ratio); + } + }(); auto const mem_stats_logger = cudf::memory_stats_logger(); state.set_cuda_stream(nvbench::make_cuda_stream_view(cudf::get_default_stream().value())); - state.exec(nvbench::exec_tag::timer | nvbench::exec_tag::sync, - [&](nvbench::launch&, auto& timer) { - cuio_source_sink_pair source_sink(io_type::FILEPATH); - - // Lift the byte caps well above the row caps so page boundaries - // are driven purely by `max_page_size_rows` (100K rows = 800KB - // for INT64, which exceeds the default 512KB page byte cap). - // This guarantees exactly `pages_per_chunk` pages per chunk, - // matching the host-side data layout. - constexpr std::size_t page_bytes_cap = std::size_t{64} << 20; - - timer.start(); - auto const write_opts = - cudf::io::parquet_writer_options::builder(source_sink.make_sink_info(), view) - .compression(cudf::io::compression_type::NONE) - .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) - .row_group_size_rows(num_rows) - .max_page_size_rows(page_size_rows) - .max_page_size_bytes(page_bytes_cap) - .build(); - cudf::io::write_parquet(write_opts); - timer.stop(); - - encoded_file_size = source_sink.size(); - }); - - state.add_element_count(static_cast(view.num_rows()), "rows"); + state.exec( + nvbench::exec_tag::timer | nvbench::exec_tag::sync, [&](nvbench::launch&, auto& timer) { + timer.start(); + auto const write_opts = + cudf::io::parquet_writer_options::builder(source_sink.make_sink_info(), table->view()) + .compression(cudf::io::compression_type::NONE) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .row_group_size_rows(num_rows) + .max_page_size_rows(page_size_rows) + .max_page_size_bytes(std::size_t{64} << 20) + .build(); + cudf::io::write_parquet(write_opts); + timer.stop(); + }); + + state.add_element_count(static_cast(table->num_rows()), "rows"); state.add_buffer_size( mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage"); - state.add_buffer_size(encoded_file_size, "encoded_file_size", "encoded_file_size"); - - // Emit the per-page dictionary RLE bit-width distribution. We do an extra - // untimed write with `STATISTICS_COLUMN` so that the OffsetIndex (which - // gives us each page's byte range) is actually written; the timed write - // above intentionally uses the default stats granularity to keep the - // benchmark's encoding path unperturbed. For this benchmark workload - // (flat INT64, no nulls, V1 headers) the first byte of each data page's - // payload is the RLE bit-width used to encode dict indices. - cuio_source_sink_pair inspect_sink(io_type::FILEPATH); - { - auto const inspect_opts = - cudf::io::parquet_writer_options::builder(inspect_sink.make_sink_info(), view) - .compression(cudf::io::compression_type::NONE) - .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) - .row_group_size_rows(num_rows) - .max_page_size_rows(page_size_rows) - .max_page_size_bytes(std::size_t{64} << 20) - .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) - .build(); - cudf::io::write_parquet(inspect_opts); - } + state.add_buffer_size(source_sink.size(), "encoded_file_size", "encoded_file_size"); - // Use the hybrid scan reader to build a `FileMetaData` with a fully - // materialized OffsetIndex for all column chunks (works regardless of - // column type, unlike `read_parquet_footers` which skips the page index - // when no string columns are present). - auto inspect_sources = cudf::io::make_datasources(inspect_sink.make_source_info()); - auto& inspect_ds = *inspect_sources.front(); - auto const footer_buf = cudf::io::parquet::fetch_footer_to_host(inspect_ds); - cudf::io::parquet::experimental::hybrid_scan_reader reader(*footer_buf, - cudf::io::parquet_reader_options{}); - auto const page_index_range = reader.page_index_byte_range(); - if (not page_index_range.is_empty()) { - auto const pi_buf = - cudf::io::parquet::fetch_page_index_to_host(inspect_ds, page_index_range); - reader.setup_page_index(*pi_buf); - } - auto const footer = reader.parquet_metadata(); - auto const page_bits = extract_page_dict_bits(inspect_ds, footer); - - if (not page_bits.empty()) { - auto const [min_it, max_it] = std::minmax_element(page_bits.begin(), page_bits.end()); - auto const sum = - std::accumulate(page_bits.begin(), page_bits.end(), std::uint64_t{0}); - auto const mean = static_cast(sum) / static_cast(page_bits.size()); + // Use hybrid scan reader to get footer with page index. + { + auto const datasource = + std::move(cudf::io::make_datasources(source_sink.make_source_info()).front()); + auto const datasource_ref = std::ref(*datasource); + auto const footer_buf = cudf::io::parquet::fetch_footer_to_host(datasource_ref); + cudf::io::parquet::experimental::hybrid_scan_reader reader(*footer_buf, + cudf::io::parquet_reader_options{}); + auto const page_index_bytes = reader.page_index_byte_range(); + CUDF_EXPECTS(not page_index_bytes.is_empty(), "Page index is required"); + auto const page_index_buffer = + cudf::io::parquet::fetch_page_index_to_host(datasource_ref, page_index_bytes); + reader.setup_page_index(*page_index_buffer); + + auto const metadata = reader.parquet_metadata(); + auto const page_rle_bits = compute_page_dict_bits(datasource_ref, metadata); + + CUDF_EXPECTS(not page_rle_bits.empty(), "No dictionary-encoded pages found"); + + auto const [min_it, max_it] = std::minmax_element(page_rle_bits.begin(), page_rle_bits.end()); + auto const sum = std::accumulate(page_rle_bits.begin(), page_rle_bits.end(), std::uint64_t{0}); + auto const mean = static_cast(sum) / static_cast(page_rle_bits.size()); state.add_element_count(static_cast(*min_it), "dict_rle_bits_min"); state.add_element_count(static_cast(*max_it), "dict_rle_bits_max"); state.add_element_count(mean, "dict_rle_bits_mean"); @@ -260,5 +228,9 @@ void BM_parq_write_dict_encoding(nvbench::state& state) NVBENCH_BENCH(BM_parq_write_dict_encoding) .set_name("parquet_write_dict_encoding") - .set_min_samples(3) - .add_int64_axis("num_rows", {1'000'000}); + .set_min_samples(4) + .add_int64_axis("reverse_order", {false, true}) + .add_int64_axis("num_rows", {1'000'000}) + .add_int64_axis("page_size_rows", {10'000, 100'000}) + .add_int64_axis("cardinality", {64'000, 100'000}) + .add_float64_axis("freq_set_ratio", {0.001, 0.01}); \ No newline at end of file From 7dd5460f9e55a50a967258ebc0dc706d45d97341 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Thu, 23 Apr 2026 21:58:52 +0000 Subject: [PATCH 05/28] Improve tests --- cpp/tests/io/parquet_misc_test.cpp | 24 ++++------- cpp/tests/io/parquet_writer_test.cpp | 61 +++++++++++++--------------- 2 files changed, 37 insertions(+), 48 deletions(-) diff --git a/cpp/tests/io/parquet_misc_test.cpp b/cpp/tests/io/parquet_misc_test.cpp index 0de0628ed0a1..6b91b1575267 100644 --- a/cpp/tests/io/parquet_misc_test.cpp +++ b/cpp/tests/io/parquet_misc_test.cpp @@ -169,21 +169,15 @@ TEST_P(ParquetSizedTest, DictionaryTest) }); EXPECT_TRUE(used_dict); - // and check that the correct number of bits was used. Phase 2 variable-bit-width - // encoding lets each page emit `ceil(log2(page_max_dict_index + 1))` bits rather - // than the chunk-wide maximum, so individual pages may use *fewer* bits than - // `GetParam()`. The chunk-wide maximum must still be reached by at least one page - // (the one that references the last-assigned dict_id), so we verify: - // - every page's bit width is <= GetParam() (upper bound is respected), and - // - max page bit width == GetParam() (the chunk-wide bound is tight). - auto const oi = read_offset_index(source, fmd.row_groups.front().columns.front()); - int max_nbits = 0; - for (auto const& pl : oi.page_locations) { - auto const nbits = read_dict_bits(source, pl); - EXPECT_LE(nbits, GetParam()); - max_nbits = std::max(max_nbits, nbits); - } - EXPECT_EQ(max_nbits, GetParam()); + // and check that the correct number of bits was used + auto const oi = read_offset_index(source, fmd.row_groups.front().columns.front()); + auto const page_with_max_bits = std::max_element( + oi.page_locations.begin(), oi.page_locations.end(), [&source](auto const& a, auto const& b) { + return read_dict_bits(source, a) < read_dict_bits(source, b); + }); + + EXPECT_NE(page_with_max_bits, oi.page_locations.end()); + EXPECT_EQ(read_dict_bits(source, *page_with_max_bits), GetParam()); } /////////////////////// diff --git a/cpp/tests/io/parquet_writer_test.cpp b/cpp/tests/io/parquet_writer_test.cpp index 15a40286d828..5bfede6c931e 100644 --- a/cpp/tests/io/parquet_writer_test.cpp +++ b/cpp/tests/io/parquet_writer_test.cpp @@ -28,6 +28,7 @@ #include +#include #include #include #include @@ -1110,6 +1111,24 @@ TEST_F(ParquetWriterTest, SingleValueDictionaryTest) EXPECT_EQ(nbits, expected_bits); } +namespace { + +/** @brief Returns whether a column chunk was actually written with dictionary encoding. + * + * @param chunk The column chunk to check. + * @return Whether the column chunk was actually written with dictionary encoding + */ +[[nodiscard]] bool is_chunk_dict_encoded(cudf::io::parquet::ColumnChunk const& chunk) +{ + return std::any_of( + chunk.meta_data.encodings.begin(), chunk.meta_data.encodings.end(), [](auto const encoding) { + return encoding == cudf::io::parquet::Encoding::PLAIN_DICTIONARY or + encoding == cudf::io::parquet::Encoding::RLE_DICTIONARY; + }); +} + +} // namespace + // Phase 2 per-page variable-bit-width RLE coverage (PHASE_2_VARIABLE_BITS.md §2.4). // The four tests below exercise the new code path from three angles: // (1) decoder correctness across cardinalities and physical types @@ -1118,31 +1137,6 @@ TEST_F(ParquetWriterTest, SingleValueDictionaryTest) // (3) the optimization actually fires -- file size drops below the pre-Phase-2 // chunk-wide number on the benchmark workload. -// `ceil(log2(max_dict_index + 1))`, floored at 1. Matches what -// `compute_page_dict_rle_bits_kernel` writes device-side (`32 - -// countl_zero(max)`, floored at 1) and the chunk-wide convention in -// `build_chunk_dictionaries` (all-null pages still emit a 1-bit RLE preamble, -// see writer_impl.cu's `std::max(..., 1)`). -namespace { -[[nodiscard]] int num_required_bits(uint32_t v) { return std::max(std::bit_width(v), 1); } - -// Returns true iff the column chunk was actually written with dictionary -// encoding. The writer may silently fall back to PLAIN for sparse-cardinality -// inputs even under `dictionary_policy::ALWAYS` (e.g., when the dict page -// would dwarf the data page), so tests that inspect per-page bit widths must -// check this before dereferencing `read_dict_bits` output. -[[nodiscard]] bool chunk_used_dictionary(cudf::io::parquet::ColumnChunk const& chunk) -{ - for (auto const enc : chunk.meta_data.encodings) { - if (enc == cudf::io::parquet::Encoding::PLAIN_DICTIONARY or - enc == cudf::io::parquet::Encoding::RLE_DICTIONARY) { - return true; - } - } - return false; -} -} // namespace - // Round-trip INT64 dictionary-encoded columns at three cardinalities. The // cardinalities are chosen to straddle the 10-bit / 16-bit / 20-bit RLE widths // so the reader has to correctly consume per-page preambles of different @@ -1189,9 +1183,9 @@ TEST_F(ParquetWriterTest, VariableBitWidthRoundTripIntegers) // when the dict page would exceed the row data (very sparse cardinality // relative to row count). Treat that as out-of-scope for bit-width // checking but still validate round-trip correctness above. - if (not chunk_used_dictionary(chunk)) { continue; } + if (not is_chunk_dict_encoded(chunk)) { continue; } - auto const chunk_wide_max_bits = num_required_bits(cardinality - 1); + auto const chunk_wide_max_bits = std::bit_width(cardinality - 1); auto const oi = read_offset_index(source, chunk); ASSERT_GT(oi.page_locations.size(), 1u) << "cardinality=" << cardinality; for (auto const& pl : oi.page_locations) { @@ -1240,7 +1234,8 @@ TEST_F(ParquetWriterTest, VariableBitWidthRoundTripStrings) cudf::io::parquet::FileMetaData fmd; read_footer(source, &fmd); - auto const chunk_wide_max_bits = num_required_bits(cardinality - 1); + auto const chunk_wide_max_bits = + std::max(std::bit_width(static_cast(cardinality - 1)), 1); auto const oi = read_offset_index(source, fmd.row_groups.front().columns.front()); ASSERT_GT(oi.page_locations.size(), 1u); for (auto const& pl : oi.page_locations) { @@ -1312,14 +1307,14 @@ TEST_F(ParquetWriterTest, VariableBitWidthRoundTripLists) auto const source = cudf::io::datasource::create(cudf::host_span{buffer_span}); cudf::io::parquet::FileMetaData fmd; read_footer(source, &fmd); - EXPECT_TRUE(chunk_used_dictionary(fmd.row_groups.front().columns.front())); + EXPECT_TRUE(is_chunk_dict_encoded(fmd.row_groups.front().columns.front())); } // End-to-end file-size gate. Reproduces the `parquet_write_dict_encoding` // benchmark workload (8 "common" pages touching only 64 frequent values + 2 // "rare" pages touching values 64..63999) and asserts: // (a) common pages bit-pack at their page-local minimum width (`nbits <= -// num_required_bits(frequent_set_size - 1)`), +// std::max(std::bit_width(frequent_set_size - 1), 1)`), // (b) at least one page reaches the chunk-wide width (so dict_rle_bits is // still the upper bound, not silently clobbered), // (c) the resulting file is strictly smaller than the Phase 1 baseline @@ -1371,15 +1366,15 @@ TEST_F(ParquetWriterTest, VariableBitWidthSmallerThanChunkWide) cudf::io::parquet::FileMetaData fmd; read_footer(source, &fmd); - ASSERT_TRUE(chunk_used_dictionary(fmd.row_groups.front().columns.front())); + ASSERT_TRUE(is_chunk_dict_encoded(fmd.row_groups.front().columns.front())); // Chunk-wide upper bound computed from the full cardinality. The writer's // `dict_rle_bits` is `NumRequiredBits(num_dict_entries - 1)`, and under // uniform sampling of `[0, cardinality)` the actual number of distinct // values drawn may land a bit below cardinality -- so we use this value as // a *ceiling* rather than an equality target. - auto const chunk_wide_max_bits = num_required_bits(cardinality - 1); - auto const frequent_max_bits = num_required_bits(frequent_set_size - 1); // 6 + auto const chunk_wide_max_bits = std::bit_width(cardinality - 1); + auto const frequent_max_bits = std::bit_width(frequent_set_size - 1); // 6 auto const oi = read_offset_index(source, fmd.row_groups.front().columns.front()); ASSERT_EQ(oi.page_locations.size(), static_cast(pages_per_chunk)); From 4600d21bbbda5a1ca0d83fceae40bf6af32e4d98 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Thu, 23 Apr 2026 22:42:43 +0000 Subject: [PATCH 06/28] Minor --- cpp/benchmarks/io/parquet/parquet_writer_dict.cpp | 6 +++--- cpp/src/io/parquet/chunk_dict.cu | 8 -------- cpp/src/io/parquet/parquet_gpu.cuh | 9 +++++++++ cpp/src/io/parquet/writer_impl.cu | 4 ++++ 4 files changed, 16 insertions(+), 11 deletions(-) diff --git a/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp b/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp index f68b51565bf1..f8019e3c2b94 100644 --- a/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp +++ b/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp @@ -57,7 +57,7 @@ std::vector build_numeric_column(cudf::size_type num_rows, "frequent_set_ratio must be between 0.0 and 1.0"); CUDF_EXPECTS(num_rows % page_size_rows == 0, "num_rows must be a multiple of page_size_rows"); static_assert(frequent_pages_ratio > 0.0 and frequent_pages_ratio < 1.0, - "hot_pages_ratio must be between 0.0 and 1.0"); + "frequent_pages_ratio must be between 0.0 and 1.0"); auto const total_pages = num_rows / page_size_rows; auto const frequent_set_threshold = @@ -159,7 +159,7 @@ void BM_parq_write_dict_encoding(nvbench::state& state) auto const num_rows = static_cast(state.get_int64("num_rows")); auto const reverse_order = static_cast(state.get_int64("reverse_order")); auto const cardinality = static_cast(state.get_int64("cardinality")); - auto const frequent_set_ratio = static_cast(state.get_float64("frequent_set_ratio")); + auto const frequent_set_ratio = static_cast(state.get_float64("freq_set_ratio")); auto const page_size_rows = static_cast(state.get_int64("page_size_rows")); CUDF_EXPECTS(page_size_rows <= num_rows and num_rows % page_size_rows == 0, @@ -233,4 +233,4 @@ NVBENCH_BENCH(BM_parq_write_dict_encoding) .add_int64_axis("num_rows", {1'000'000}) .add_int64_axis("page_size_rows", {10'000, 100'000}) .add_int64_axis("cardinality", {64'000, 100'000}) - .add_float64_axis("freq_set_ratio", {0.001, 0.01}); \ No newline at end of file + .add_float64_axis("freq_set_ratio", {0.001, 0.01}); diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index 33c65a521b3f..b0849f3720b2 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -24,14 +24,6 @@ namespace cudf::io::parquet::detail { namespace { constexpr int DEFAULT_BLOCK_SIZE = 256; - -// Upper bound on the number of fragments per column chunk that the -// shared-memory histogram in `collect_map_entries_kernel` can accommodate. -// A typical workload is 1M row groups / ~5000-row fragments ≈ 200 fragments -// per chunk, so 1024 is a comfortable ceiling. If a future workload exceeds -// this, `collect_map_entries_kernel` will trip its `cudf_assert` below and -// we should add a global-memory fallback rather than silently truncating. -constexpr size_type MAX_FRAGMENTS_PER_BLOCK = 1024; } // namespace template diff --git a/cpp/src/io/parquet/parquet_gpu.cuh b/cpp/src/io/parquet/parquet_gpu.cuh index cbd1b4cf91a9..3a963adccf31 100644 --- a/cpp/src/io/parquet/parquet_gpu.cuh +++ b/cpp/src/io/parquet/parquet_gpu.cuh @@ -30,6 +30,15 @@ auto constexpr bucket_size = auto constexpr occupancy_factor = 1.43f; ///< cuCollections suggests using a hash map of size ///< N * (1/0.7) = 1.43 to target a 70% occupancy factor. +// Upper bound on the number of fragments per column chunk that the +// shared-memory histogram in `collect_map_entries_kernel` can accommodate. +// A typical workload is 1M row groups / ~5000-row fragments ≈ 200 fragments +// per chunk, so 1024 is a comfortable ceiling. Host-side code must enforce +// this before launching the kernel (see `build_chunk_dictionaries`); the +// kernel also has a `cudf_assert` as a debug-build safety net, but that is +// compiled out in release builds. +constexpr size_type MAX_FRAGMENTS_PER_BLOCK = 1024; + auto constexpr KEY_SENTINEL = key_type{-1}; auto constexpr VALUE_SENTINEL = mapped_type{-1}; auto constexpr SCOPE = cuda::thread_scope_block; diff --git a/cpp/src/io/parquet/writer_impl.cu b/cpp/src/io/parquet/writer_impl.cu index e9e95a0828cf..faa358d805ad 100644 --- a/cpp/src/io/parquet/writer_impl.cu +++ b/cpp/src/io/parquet/writer_impl.cu @@ -1885,6 +1885,10 @@ auto convert_table_to_parquet_data(table_input_metadata& table_meta, auto& row_group = agg_meta->file(p).row_groups[global_r]; auto const fragments_in_chunk = util::div_rounding_up_safe(row_group.num_rows, max_page_fragment_size); + CUDF_EXPECTS(fragments_in_chunk <= static_cast(MAX_FRAGMENTS_PER_BLOCK), + "Number of fragments per column chunk exceeds the maximum supported by the " + "parquet writer's dictionary collection kernel. Consider increasing the " + "fragment size or reducing the row group size."); row_group.total_byte_size = 0; row_group.columns.resize(num_columns); for (int c = 0; c < num_columns; c++) { From 14469b95ecb36bf58d76f7f0f099406d24bdc591 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Mon, 27 Apr 2026 18:38:35 +0000 Subject: [PATCH 07/28] Minor --- .../io/parquet/parquet_writer_dict.cpp | 14 +- cpp/src/io/parquet/chunk_dict.cu | 17 ++ cpp/tests/io/parquet_misc_test.cpp | 22 +- cpp/tests/io/parquet_writer_test.cpp | 193 +++++++++--------- 4 files changed, 134 insertions(+), 112 deletions(-) diff --git a/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp b/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp index f8019e3c2b94..bbb967a4d1bf 100644 --- a/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp +++ b/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp @@ -212,14 +212,16 @@ void BM_parq_write_dict_encoding(nvbench::state& state) cudf::io::parquet::fetch_page_index_to_host(datasource_ref, page_index_bytes); reader.setup_page_index(*page_index_buffer); - auto const metadata = reader.parquet_metadata(); - auto const page_rle_bits = compute_page_dict_bits(datasource_ref, metadata); + auto const metadata = reader.parquet_metadata(); + auto const page_dict_bits = compute_page_dict_bits(datasource_ref, metadata); - CUDF_EXPECTS(not page_rle_bits.empty(), "No dictionary-encoded pages found"); + CUDF_EXPECTS(not page_dict_bits.empty(), "No dictionary-encoded pages found"); - auto const [min_it, max_it] = std::minmax_element(page_rle_bits.begin(), page_rle_bits.end()); - auto const sum = std::accumulate(page_rle_bits.begin(), page_rle_bits.end(), std::uint64_t{0}); - auto const mean = static_cast(sum) / static_cast(page_rle_bits.size()); + auto const [min_it, max_it] = std::minmax_element(page_dict_bits.begin(), page_dict_bits.end()); + auto const sum = + std::accumulate(page_dict_bits.begin(), page_dict_bits.end(), std::uint64_t{0}); + auto const mean = + std::round(static_cast(sum) / static_cast(page_dict_bits.size())); state.add_element_count(static_cast(*min_it), "dict_rle_bits_min"); state.add_element_count(static_cast(*max_it), "dict_rle_bits_max"); state.add_element_count(mean, "dict_rle_bits_mean"); diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index b0849f3720b2..b6684d2a8172 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -277,6 +277,23 @@ CUDF_KERNEL void __launch_bounds__(block_size) // the alternative of launching a separate `cub::DeviceHistogram` + // `DeviceScan` per chunk (kernel-launch overhead dominates for many small // chunks; see PROBLEM.md §5.1 and PHASE_1_ORDERING.md §1.2.2). +// +// Performance note: the inner shared-atomic `atomicAdd` in Pass 1/3 is *not* +// the bottleneck for this kernel on sm_100. Profiling (see +// `l1tex__data_bank_conflicts_pipe_lsu.sum` + warp-state stalls) shows the +// kernel is stall-bound on L1TEX scoreboard dependencies, i.e. the global +// loads of `chunk_slots[slot_idx]`, with SMs mostly idle because the grid is +// `num_chunks` blocks (often << number of SMs). Splitting the histogram into +// per-warp-group replicas to reduce atomic contention was tried and produced +// no measurable speedup (same 1.04 ms kernel time on B200 at 32-way vs 8-way +// replica contention). Future optimization effort should target either: +// - Wider parallelism (one block per N fragments instead of per chunk), or +// - Coalesced prefetch of the slot array into shared memory per pass to +// amortize the global-load stalls across Pass 1 and Pass 3. +// The `cuda::atomic_ref<..., thread_scope_block>::fetch_add(..., relaxed)` +// alternative was also tried; on CUDA 13.1 it lowers to the generic LSU +// atomic pipe (`ATOM.E.ADD.S32.STRONG.SM`) instead of the dedicated shared- +// atomic unit (`ATOMS.ADD`), adding ~75% to kernel runtime. template CUDF_KERNEL void __launch_bounds__(block_size) collect_map_entries_kernel(device_span const map_storage, diff --git a/cpp/tests/io/parquet_misc_test.cpp b/cpp/tests/io/parquet_misc_test.cpp index 6b91b1575267..67803f1592d4 100644 --- a/cpp/tests/io/parquet_misc_test.cpp +++ b/cpp/tests/io/parquet_misc_test.cpp @@ -12,7 +12,9 @@ #include +#include #include +#include //////////////////////////////// // delta encoding writer tests @@ -170,14 +172,18 @@ TEST_P(ParquetSizedTest, DictionaryTest) EXPECT_TRUE(used_dict); // and check that the correct number of bits was used - auto const oi = read_offset_index(source, fmd.row_groups.front().columns.front()); - auto const page_with_max_bits = std::max_element( - oi.page_locations.begin(), oi.page_locations.end(), [&source](auto const& a, auto const& b) { - return read_dict_bits(source, a) < read_dict_bits(source, b); - }); - - EXPECT_NE(page_with_max_bits, oi.page_locations.end()); - EXPECT_EQ(read_dict_bits(source, *page_with_max_bits), GetParam()); + auto const oi = read_offset_index(source, fmd.row_groups.front().columns.front()); + std::vector page_dict_bits; + page_dict_bits.reserve(oi.page_locations.size()); + std::transform( + oi.page_locations.begin(), + oi.page_locations.end(), + std::back_inserter(page_dict_bits), + [&source](auto const& page_location) { return read_dict_bits(source, page_location); }); + + auto const max_bits = std::max_element(page_dict_bits.begin(), page_dict_bits.end()); + EXPECT_NE(max_bits, page_dict_bits.end()); + EXPECT_EQ(*max_bits, GetParam()); } /////////////////////// diff --git a/cpp/tests/io/parquet_writer_test.cpp b/cpp/tests/io/parquet_writer_test.cpp index 5bfede6c931e..41a1c0311696 100644 --- a/cpp/tests/io/parquet_writer_test.cpp +++ b/cpp/tests/io/parquet_writer_test.cpp @@ -33,6 +33,7 @@ #include #include #include +#include #include using cudf::test::iterators::no_nulls; @@ -1127,6 +1128,51 @@ namespace { }); } +[[nodiscard]] std::vector read_page_dict_bits( + std::unique_ptr const& source, cudf::io::parquet::ColumnChunk const& chunk) +{ + auto const oi = read_offset_index(source, chunk); + std::vector page_dict_bits; + page_dict_bits.reserve(oi.page_locations.size()); + std::transform( + oi.page_locations.begin(), + oi.page_locations.end(), + std::back_inserter(page_dict_bits), + [&source](auto const& page_location) { return read_dict_bits(source, page_location); }); + return page_dict_bits; +} + +template +[[nodiscard]] auto write_and_read_dict_bits(table_view const& expected, + ConfigureWriter&& configure_writer) + -> std::pair, cudf::io::parquet::FileMetaData> +{ + auto buffer = std::vector{}; + auto builder = cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&buffer}, expected) + .compression(cudf::io::compression_type::NONE) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS); + configure_writer(builder); + + cudf::io::parquet_writer_options out_opts = builder; + cudf::io::write_parquet(out_opts); + + auto const buffer_span = + cudf::host_span(reinterpret_cast(buffer.data()), buffer.size()); + cudf::io::parquet_reader_options in_opts = + cudf::io::parquet_reader_options::builder(cudf::io::source_info(buffer_span)); + auto const result = cudf::io::read_parquet(in_opts); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); + + auto source = cudf::io::datasource::create(cudf::host_span{buffer_span}); + cudf::io::parquet::FileMetaData fmd; + read_footer(source, &fmd); + + auto const& chunk = fmd.row_groups.front().columns.front(); + return {is_chunk_dict_encoded(chunk) ? read_page_dict_bits(source, chunk) : std::vector{}, + std::move(fmd)}; +} + } // namespace // Phase 2 per-page variable-bit-width RLE coverage (PHASE_2_VARIABLE_BITS.md §2.4). @@ -1155,28 +1201,10 @@ TEST_F(ParquetWriterTest, VariableBitWidthRoundTripIntegers) std::generate(values.begin(), values.end(), [&] { return dist(rng); }); auto const col = cudf::test::fixed_width_column_wrapper(values.begin(), values.end()); - auto const expected = table_view{{col}}; - - auto buffer = std::vector{}; - cudf::io::parquet_writer_options out_opts = - cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&buffer}, expected) - .compression(cudf::io::compression_type::NONE) - .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) - .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) - .row_group_size_rows(nrows) - .max_page_size_rows(nrows / 4); - cudf::io::write_parquet(out_opts); - - auto const buffer_span = - cudf::host_span(reinterpret_cast(buffer.data()), buffer.size()); - cudf::io::parquet_reader_options in_opts = - cudf::io::parquet_reader_options::builder(cudf::io::source_info(buffer_span)); - auto const result = cudf::io::read_parquet(in_opts); - CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); - - auto const source = cudf::io::datasource::create(cudf::host_span{buffer_span}); - cudf::io::parquet::FileMetaData fmd; - read_footer(source, &fmd); + auto const expected = table_view{{col}}; + auto const [page_dict_bits, fmd] = write_and_read_dict_bits(expected, [=](auto& builder) { + builder.row_group_size_rows(nrows).max_page_size_rows(nrows / 4); + }); auto const& chunk = fmd.row_groups.front().columns.front(); // Under `dictionary_policy::ALWAYS` the writer still falls back to PLAIN @@ -1186,10 +1214,8 @@ TEST_F(ParquetWriterTest, VariableBitWidthRoundTripIntegers) if (not is_chunk_dict_encoded(chunk)) { continue; } auto const chunk_wide_max_bits = std::bit_width(cardinality - 1); - auto const oi = read_offset_index(source, chunk); - ASSERT_GT(oi.page_locations.size(), 1u) << "cardinality=" << cardinality; - for (auto const& pl : oi.page_locations) { - auto const nbits = read_dict_bits(source, pl); + ASSERT_GT(page_dict_bits.size(), 1u) << "cardinality=" << cardinality; + for (auto const nbits : page_dict_bits) { EXPECT_GE(nbits, 1) << "cardinality=" << cardinality; EXPECT_LE(nbits, chunk_wide_max_bits) << "cardinality=" << cardinality; } @@ -1212,34 +1238,15 @@ TEST_F(ParquetWriterTest, VariableBitWidthRoundTripStrings) auto const col = cudf::test::strings_column_wrapper(values.begin(), values.end()); auto const expected = table_view{{col}}; - - auto buffer = std::vector{}; - cudf::io::parquet_writer_options out_opts = - cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&buffer}, expected) - .compression(cudf::io::compression_type::NONE) - .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) - .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) - .row_group_size_rows(nrows) - .max_page_size_rows(nrows / 4); - cudf::io::write_parquet(out_opts); - - auto const buffer_span = - cudf::host_span(reinterpret_cast(buffer.data()), buffer.size()); - cudf::io::parquet_reader_options in_opts = - cudf::io::parquet_reader_options::builder(cudf::io::source_info(buffer_span)); - auto const result = cudf::io::read_parquet(in_opts); - CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); - - auto const source = cudf::io::datasource::create(cudf::host_span{buffer_span}); - cudf::io::parquet::FileMetaData fmd; - read_footer(source, &fmd); + auto const [page_dict_bits, fmd] = write_and_read_dict_bits(expected, [=](auto& builder) { + builder.row_group_size_rows(nrows).max_page_size_rows(nrows / 4); + }); auto const chunk_wide_max_bits = std::max(std::bit_width(static_cast(cardinality - 1)), 1); - auto const oi = read_offset_index(source, fmd.row_groups.front().columns.front()); - ASSERT_GT(oi.page_locations.size(), 1u); - for (auto const& pl : oi.page_locations) { - auto const nbits = read_dict_bits(source, pl); + ASSERT_TRUE(is_chunk_dict_encoded(fmd.row_groups.front().columns.front())); + ASSERT_GT(page_dict_bits.size(), 1u); + for (auto const nbits : page_dict_bits) { EXPECT_GE(nbits, 1); EXPECT_LE(nbits, chunk_wide_max_bits); } @@ -1285,28 +1292,11 @@ TEST_F(ParquetWriterTest, VariableBitWidthRoundTripLists) .release(); auto list_col = cudf::make_lists_column( num_lists, std::move(offsets_col), std::move(leaf_col), 0, rmm::device_buffer{}); - auto const expected = table_view{{*list_col}}; - - auto buffer = std::vector{}; - cudf::io::parquet_writer_options out_opts = - cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&buffer}, expected) - .compression(cudf::io::compression_type::NONE) - .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) - .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) - .row_group_size_rows(num_lists) - .max_page_size_rows(num_lists / 4); - cudf::io::write_parquet(out_opts); - - auto const buffer_span = - cudf::host_span(reinterpret_cast(buffer.data()), buffer.size()); - cudf::io::parquet_reader_options in_opts = - cudf::io::parquet_reader_options::builder(cudf::io::source_info(buffer_span)); - auto const result = cudf::io::read_parquet(in_opts); - CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); - - auto const source = cudf::io::datasource::create(cudf::host_span{buffer_span}); - cudf::io::parquet::FileMetaData fmd; - read_footer(source, &fmd); + auto const expected = table_view{{*list_col}}; + auto const [page_dict_bits, fmd] = write_and_read_dict_bits(expected, [=](auto& builder) { + builder.row_group_size_rows(num_lists).max_page_size_rows(num_lists / 4); + }); + static_cast(page_dict_bits); EXPECT_TRUE(is_chunk_dict_encoded(fmd.row_groups.front().columns.front())); } @@ -1342,31 +1332,40 @@ TEST_F(ParquetWriterTest, VariableBitWidthSmallerThanChunkWide) } auto const col = cudf::test::fixed_width_column_wrapper(values.begin(), values.end()); - auto const expected = table_view{{col}}; + auto const expected = table_view{{col}}; + auto buffer = std::vector{}; + auto const [page_dict_bits, fmd] = [&]() { + auto local_buffer = std::vector{}; + auto builder = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&local_buffer}, expected) + .compression(cudf::io::compression_type::NONE) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) + .row_group_size_rows(num_rows) + .max_page_size_rows(page_size_rows) + .max_page_size_bytes(std::size_t{64} << 20); + cudf::io::parquet_writer_options out_opts = builder; + cudf::io::write_parquet(out_opts); - auto buffer = std::vector{}; - cudf::io::parquet_writer_options out_opts = - cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&buffer}, expected) - .compression(cudf::io::compression_type::NONE) - .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) - .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) - .row_group_size_rows(num_rows) - .max_page_size_rows(page_size_rows) - .max_page_size_bytes(std::size_t{64} << 20); - cudf::io::write_parquet(out_opts); + auto const buffer_span = cudf::host_span( + reinterpret_cast(local_buffer.data()), local_buffer.size()); + cudf::io::parquet_reader_options in_opts = + cudf::io::parquet_reader_options::builder(cudf::io::source_info(buffer_span)); + auto const result = cudf::io::read_parquet(in_opts); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); - auto const buffer_span = - cudf::host_span(reinterpret_cast(buffer.data()), buffer.size()); - cudf::io::parquet_reader_options in_opts = - cudf::io::parquet_reader_options::builder(cudf::io::source_info(buffer_span)); - auto const result = cudf::io::read_parquet(in_opts); - CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); + auto source = cudf::io::datasource::create(cudf::host_span{buffer_span}); + cudf::io::parquet::FileMetaData local_fmd; + read_footer(source, &local_fmd); - auto const source = cudf::io::datasource::create(cudf::host_span{buffer_span}); - cudf::io::parquet::FileMetaData fmd; - read_footer(source, &fmd); + buffer = std::move(local_buffer); + auto const& chunk = local_fmd.row_groups.front().columns.front(); + return std::pair{ + is_chunk_dict_encoded(chunk) ? read_page_dict_bits(source, chunk) : std::vector{}, + std::move(local_fmd)}; + }(); - ASSERT_TRUE(is_chunk_dict_encoded(fmd.row_groups.front().columns.front())); + EXPECT_TRUE(is_chunk_dict_encoded(fmd.row_groups.front().columns.front())); // Chunk-wide upper bound computed from the full cardinality. The writer's // `dict_rle_bits` is `NumRequiredBits(num_dict_entries - 1)`, and under @@ -1376,14 +1375,12 @@ TEST_F(ParquetWriterTest, VariableBitWidthSmallerThanChunkWide) auto const chunk_wide_max_bits = std::bit_width(cardinality - 1); auto const frequent_max_bits = std::bit_width(frequent_set_size - 1); // 6 - auto const oi = read_offset_index(source, fmd.row_groups.front().columns.front()); - ASSERT_EQ(oi.page_locations.size(), static_cast(pages_per_chunk)); + ASSERT_EQ(page_dict_bits.size(), static_cast(pages_per_chunk)); int common_page_count = 0; int rare_page_count = 0; int max_observed_bits = 0; - for (auto const& pl : oi.page_locations) { - auto const nbits = read_dict_bits(source, pl); + for (auto const nbits : page_dict_bits) { EXPECT_GE(nbits, 1); EXPECT_LE(nbits, chunk_wide_max_bits); max_observed_bits = std::max(max_observed_bits, nbits); From fda8a8a89eaf8cc3ec7cb602bccd49aa13e3e202 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 01:11:56 +0000 Subject: [PATCH 08/28] Humanize phase 1 --- cpp/src/io/parquet/chunk_dict.cu | 386 +++++++++++------------------ cpp/src/io/parquet/page_enc.cu | 2 - cpp/src/io/parquet/parquet_gpu.cuh | 11 +- cpp/src/io/parquet/parquet_gpu.hpp | 4 +- cpp/src/io/parquet/writer_impl.cu | 12 +- 5 files changed, 163 insertions(+), 252 deletions(-) diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index b6684d2a8172..81bbc5cd8747 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -23,8 +23,8 @@ namespace cudf::io::parquet::detail { namespace { + constexpr int DEFAULT_BLOCK_SIZE = 256; -} // namespace template struct equality_functor { @@ -51,6 +51,8 @@ template struct map_insert_fn { storage_ref_type const& storage_ref; EncColumnChunk* const& chunk; + PageFragment* const& frag; + mapped_type const frag_idx; template __device__ void operator()(size_type const s_start_value_idx, size_type const end_value_idx) @@ -62,6 +64,9 @@ struct map_insert_fn { auto const col = chunk->col_desc; column_device_view const& data_col = *col->leaf_column; __shared__ size_type total_num_dict_entries; + __shared__ size_type num_dict_vals; + if (threadIdx.x == 0) { num_dict_vals = 0; } + __syncthreads(); using equality_fn_type = equality_functor; using hash_fn_type = hash_functor; @@ -86,7 +91,7 @@ struct map_insert_fn { cuda::atomic_ref const chunk_uniq_data_size{chunk->uniq_data_size}; // Note: Adjust the following loop to use `cg::tile` if needed in the future. - for (thread_index_type val_idx = s_start_value_idx + t; val_idx - t < end_value_idx; + for (size_type val_idx = s_start_value_idx + t; val_idx - t < end_value_idx; val_idx += block_size) { size_type is_unique = 0; size_type uniq_elem_size = 0; @@ -98,50 +103,23 @@ struct map_insert_fn { // Insert tile_val_idx to hash map and count successful insertions. if (is_valid) { // Insert the keys using a single thread for best performance for now. - // Stamp `slot.second` with the (approximate) lowest column-relative fragment - // index that inserted this value. `cuco::static_map_ref::insert` writes the - // value exactly once per key -- on the winning insertion -- so after this - // kernel, every non-empty slot's `second` field equals the `blockIdx.x` of - // the block that first won that slot. Block scheduling on NVIDIA GPUs is - // approximately monotonic in `blockIdx`, giving a strong first-appearance - // signal that `collect_map_entries_kernel` converts into monotone dict_ids. - // - // NOTE: `blockIdx.x` here is the *column-relative* fragment index (the - // populate grid is (num_fragments_per_col, num_cols)). Each `EncColumnChunk` - // covers a contiguous range of these indices within its column, so the - // collect kernel subtracts the chunk's first column-relative fragment to get - // a chunk-local histogram bucket. - // - // The pair is constructed as `slot_type` (exactly `value_type` for the map, - // 8 bytes total) rather than the tempting `cuco::pair{val_idx, hint}`. The - // latter would deduce `cuco::pair` because `val_idx` is - // `thread_index_type` (int64_t), and cuco's `packed_cas` path (selected when - // `sizeof(value_type) <= 8`) reinterprets the input as a single `uint64_t` - // -- it would then CAS the low 8 bytes (the `int64_t` key) into the slot and - // silently drop the payload, leaving `slot.second = 0`. - // - // TODO: when the cuco pin exposes `static_map::insert_or_apply` with - // `cuco::op::min`, switch to that for exact first-fragment semantics - // (PROBLEM.md §5.2 Option II). - auto const fragment_hint = static_cast(blockIdx.x); - is_unique = - map_insert_ref.insert(slot_type{static_cast(val_idx), fragment_hint}); + is_unique = map_insert_ref.insert(slot_type{static_cast(val_idx), frag_idx}); uniq_elem_size = [&]() -> size_type { if (not is_unique) { return 0; } switch (col->physical_type) { - case Type::INT32: return 4; - case Type::INT64: return 8; - case Type::INT96: return 12; - case Type::FLOAT: return 4; - case Type::DOUBLE: return 8; + case Type::INT32: return sizeof(int32_t); + case Type::INT64: return sizeof(int64_t); + case Type::INT96: return sizeof(int32_t) + sizeof(int64_t); + case Type::FLOAT: return sizeof(float); + case Type::DOUBLE: return sizeof(double); case Type::BYTE_ARRAY: { auto const col_type = data_col.type().id(); if (col_type == type_id::STRING) { - // Strings are stored as 4 byte length + string bytes - return 4 + data_col.element(val_idx).size_bytes(); + // Strings are stored as int32_t length + string bytes + return sizeof(int32_t) + data_col.element(val_idx).size_bytes(); } else if (col_type == type_id::LIST) { - // Binary is stored as 4 byte length + bytes - return 4 + + // Binary is stored as int32_t length + bytes + return sizeof(int32_t) + get_element(data_col, val_idx).size_bytes(); } CUDF_UNREACHABLE( @@ -161,18 +139,22 @@ struct map_insert_fn { auto num_unique = block_reduce(reduce_storage).Sum(is_unique); __syncthreads(); auto uniq_data_size = block_reduce(reduce_storage).Sum(uniq_elem_size); - // The first thread in the block atomically updates total num_unique and uniq_data_size + // The first thread in the block atomically updates total num_unique and uniq_data_size, + // and accumulates the per-fragment winning-insert count. if (t == 0) { total_num_dict_entries = chunk_num_dict_entries.fetch_add(num_unique, cuda::std::memory_order_relaxed); total_num_dict_entries += num_unique; + num_dict_vals += num_unique; chunk_uniq_data_size.fetch_add(uniq_data_size, cuda::std::memory_order_relaxed); } __syncthreads(); // Check if the num unique values in chunk has already exceeded max dict size and early exit - if (total_num_dict_entries > MAX_DICT_SIZE) { return; } + if (total_num_dict_entries > MAX_DICT_SIZE) { break; } } // for loop + // Flush the number of unique values inserted to the fragment. + if (t == 0) { frag->num_dict_vals = num_dict_vals; } } else { CUDF_UNREACHABLE("Unsupported type to insert in map"); } @@ -189,8 +171,8 @@ struct map_find_fn { size_type const s_ck_start_val_idx) { if constexpr (column_device_view::has_element_accessor()) { - auto const col = chunk->col_desc; - column_device_view const& data_col = *col->leaf_column; + auto const col = chunk->col_desc; + auto const& data_col = *col->leaf_column; using equality_fn_type = equality_functor; using hash_fn_type = hash_functor; @@ -232,13 +214,13 @@ struct map_find_fn { template CUDF_KERNEL void __launch_bounds__(block_size) populate_chunk_hash_maps_kernel(device_span const map_storage, - cudf::detail::device_2dspan frags) + cudf::detail::device_2dspan frags) { - auto const col_idx = blockIdx.y; - auto const block_x = blockIdx.x; - auto const frag = frags[col_idx][block_x]; - auto chunk = frag.chunk; - auto col = chunk->col_desc; + auto const col_idx = blockIdx.y; + auto const frag_idx = blockIdx.x; + auto& frag = frags[col_idx][frag_idx]; + auto chunk = frag.chunk; + auto col = chunk->col_desc; if (not chunk->use_dictionary) { return; } @@ -252,142 +234,92 @@ CUDF_KERNEL void __launch_bounds__(block_size) column_device_view const& data_col = *col->leaf_column; storage_ref_type const storage_ref{chunk->dict_map_size, map_storage.data() + chunk->dict_map_offset}; - type_dispatcher(data_col.type(), - map_insert_fn{storage_ref, chunk}, - s_start_value_idx, - end_value_idx); + type_dispatcher( + data_col.type(), + map_insert_fn{storage_ref, chunk, &frag, static_cast(frag_idx)}, + s_start_value_idx, + end_value_idx); } -// Assigns monotone `dict_id`s bucketed by the fragment-hint that -// `populate_chunk_hash_maps_kernel` stamped into every slot's `second` field. -// After this kernel, `slot.second` has been overwritten with the final -// `dict_id` (an index into `chunk.dict_data`), and all values first inserted -// by fragment `f_i` are assigned `dict_id`s strictly less than values first -// inserted by fragment `f_j` for `f_i < f_j`. Per-page max `dict_index` is -// therefore monotone in page order whenever the underlying value distribution -// is, which is the invariant Phase 2 exploits for per-page bit-width savings. -// -// Algorithm (one block per chunk, three passes over `dict_map_size` slots): -// 1. Histogram: count non-empty slots by (fragment_hint - f_start). -// 2. BlockScan the histogram to exclusive-prefix offsets. -// 3. Claim `dict_id = atomicAdd(&fragment_cursor[bucket], 1)` per slot, -// overwrite slot.second, and write `dict_data[dict_id] = key`. -// -// Three passes of `dict_map_size / block_size` iterations are cheaper than -// the alternative of launching a separate `cub::DeviceHistogram` + -// `DeviceScan` per chunk (kernel-launch overhead dominates for many small -// chunks; see PROBLEM.md §5.1 and PHASE_1_ORDERING.md §1.2.2). -// -// Performance note: the inner shared-atomic `atomicAdd` in Pass 1/3 is *not* -// the bottleneck for this kernel on sm_100. Profiling (see -// `l1tex__data_bank_conflicts_pipe_lsu.sum` + warp-state stalls) shows the -// kernel is stall-bound on L1TEX scoreboard dependencies, i.e. the global -// loads of `chunk_slots[slot_idx]`, with SMs mostly idle because the grid is -// `num_chunks` blocks (often << number of SMs). Splitting the histogram into -// per-warp-group replicas to reduce atomic contention was tried and produced -// no measurable speedup (same 1.04 ms kernel time on B200 at 32-way vs 8-way -// replica contention). Future optimization effort should target either: -// - Wider parallelism (one block per N fragments instead of per chunk), or -// - Coalesced prefetch of the slot array into shared memory per pass to -// amortize the global-load stalls across Pass 1 and Pass 3. -// The `cuda::atomic_ref<..., thread_scope_block>::fetch_add(..., relaxed)` -// alternative was also tried; on CUDA 13.1 it lowers to the generic LSU -// atomic pipe (`ATOM.E.ADD.S32.STRONG.SM`) instead of the dedicated shared- -// atomic unit (`ATOMS.ADD`), adding ~75% to kernel runtime. template CUDF_KERNEL void __launch_bounds__(block_size) collect_map_entries_kernel(device_span const map_storage, device_span chunks, cudf::detail::device_2dspan frags) { - static_assert(block_size >= MAX_FRAGMENTS_PER_BLOCK, - "block_size must be >= MAX_FRAGMENTS_PER_BLOCK so one BlockScan thread backs " - "each histogram bucket."); - auto& chunk = chunks[blockIdx.x]; if (not chunk.use_dictionary) { return; } - auto const t = threadIdx.x; - auto const col_idx = chunk.col_desc_id; - auto const col_frag = frags[col_idx]; - - // Resolve the chunk's column-relative fragment range [f_start, f_start + n_frags). - // `chunk.fragments` points into `col_frag` (set in writer_impl.cu during - // row_group_fragments setup, before build_chunk_dictionaries runs), so the - // difference is the first fragment index stamped into any of this chunk's - // slots by populate_chunk_hash_maps_kernel. - __shared__ size_type f_start; - __shared__ size_type n_frags; - if (t == 0) { - f_start = static_cast(chunk.fragments - col_frag.data()); - size_type n = 0; - auto const total_col_frags = static_cast(col_frag.size()); - while (f_start + n < total_col_frags && col_frag[f_start + n].chunk == &chunk) { - ++n; - } - n_frags = n; - } - __syncthreads(); - - cudf_assert(n_frags > 0 && n_frags <= MAX_FRAGMENTS_PER_BLOCK && - "Fragments-per-chunk out of range for shared-memory histogram; raise " - "MAX_FRAGMENTS_PER_BLOCK or add a global-memory fallback."); - - // `fragment_count` starts as the per-bucket histogram, becomes per-bucket - // exclusive offsets after the scan, then is preserved as a reference for - // the MAX_DICT_SIZE overflow assert in Pass 3. `fragment_cursor` mirrors - // the offsets and is what threads `atomicAdd` into to claim dict_ids. - __shared__ size_type fragment_count[MAX_FRAGMENTS_PER_BLOCK]; - __shared__ size_type fragment_cursor[MAX_FRAGMENTS_PER_BLOCK]; - - using block_scan = cub::BlockScan; - __shared__ typename block_scan::TempStorage scan_storage; - - // Zero the histogram (only the live range; the tail past n_frags is - // untouched and never read). - if (t < n_frags) { fragment_count[t] = 0; } - __syncthreads(); - - auto* const chunk_slots = map_storage.data() + chunk.dict_map_offset; - auto const dict_map_size = chunk.dict_map_size; - - // Pass 1: histogram slot.second values into per-fragment buckets. - // `slot.second` is the column-relative fragment index that won the insert in - // `populate_chunk_hash_maps_kernel`; subtracting `f_start` normalizes it into - // a chunk-local bucket index `[0, n_frags)`. - for (size_type slot_idx = t; slot_idx < dict_map_size; slot_idx += block_size) { - auto const slot_key = chunk_slots[slot_idx].first; - if (slot_key != KEY_SENTINEL) { - auto const frag_local = chunk_slots[slot_idx].second - f_start; - cudf_assert(frag_local >= 0 && frag_local < n_frags && - "populate stamped a fragment hint outside this chunk's fragment range"); - atomicAdd(&fragment_count[frag_local], 1); + auto t = threadIdx.x; + + // Resolve the chunk's column-relative fragment range [frag_start, frag_start + num_frags). + // Both values come directly from host-populated fields on the chunk; no in-kernel reduction + // or shared memory is needed. `chunk.fragments` points into `col_frags` (set in + // writer_impl.cu during row_group_fragments setup, before build_chunk_dictionaries runs), + // so the subtraction is the first fragment index stamped into any of this chunk's slots + // by populate_chunk_hash_maps_kernel. `chunk.num_fragments` is the run length. + auto const num_frags = chunk.num_fragments; + + if (num_frags <= MAX_FRAGMENTS_PER_BLOCK) { + auto const col_idx = chunk.col_desc_id; + auto const& col_frags = frags[col_idx]; + auto const frag_start = static_cast(chunk.fragments - col_frags.data()); + + // Per-bucket cursors: initialized to the exclusive-prefix offsets of the + // per-fragment winning-insert counts; threads `atomicAdd` into these to + // claim dict_ids in Pass 2. + __shared__ size_type fragment_cursor[MAX_FRAGMENTS_PER_BLOCK]; + + // Pass 1: in-block exclusive scan over the per-fragment winning-insert + // counts populate wrote into each fragment. This replaces the old + // slot-rescan histogram: `frag.num_dict_vals` already equals + // "number of slots in `chunk_slots[]` stamped with column-relative fragment + // index (f_start + i)", because populate's `blockIdx.x` is exactly that + // fragment index and each winning CAS contributes exactly one stamp. + { + using block_scan = cub::BlockScan; + __shared__ typename block_scan::TempStorage scan_storage; + + auto const per_thread_count = (t < num_frags) ? col_frags[frag_start + t].num_dict_vals : 0; + auto per_thread_offset = 0; + block_scan(scan_storage).ExclusiveSum(per_thread_count, per_thread_offset); + if (t < num_frags) { fragment_cursor[t] = per_thread_offset; } } - } - __syncthreads(); - - // Pass 2: in-block exclusive scan of the histogram -> offsets -> cursors. - { - size_type const per_thread_count = (t < n_frags) ? fragment_count[t] : 0; - size_type per_thread_offset = 0; - block_scan(scan_storage).ExclusiveSum(per_thread_count, per_thread_offset); - if (t < n_frags) { - fragment_count[t] = per_thread_offset; - fragment_cursor[t] = per_thread_offset; + __syncthreads(); + + // Iterate over all slots in the map, claim a dict_id in the bucket and write it to dict_data + for (; t < chunk.dict_map_size; t += block_size) { + auto* slot = map_storage.data() + chunk.dict_map_offset + t; + auto const key = slot->first; + if (key != KEY_SENTINEL) { + auto const frag_loc = static_cast(slot->second) - frag_start; + cudf_assert(frag_loc >= 0 && frag_loc < num_frags && + "populate stamped a fragment hint outside this chunk's fragment range"); + auto const loc = atomicAdd(&fragment_cursor[frag_loc], 1); + cudf_assert(loc < MAX_DICT_SIZE && "Number of filled slots exceeds max dict size"); + chunk.dict_data[loc] = key; + slot->second = loc; + } } - } - __syncthreads(); - - // Pass 3: claim a dict_id in the bucket, overwrite slot.second, materialize - // dict_data. The atomicAdd is shared-memory only, so no global traffic. - for (size_type slot_idx = t; slot_idx < dict_map_size; slot_idx += block_size) { - auto const slot_key = chunk_slots[slot_idx].first; - if (slot_key != KEY_SENTINEL) { - auto const frag_local = chunk_slots[slot_idx].second - f_start; - auto const loc = atomicAdd(&fragment_cursor[frag_local], 1); - cudf_assert(loc < MAX_DICT_SIZE && "Number of filled slots exceeds max dict size"); - chunk.dict_data[loc] = slot_key; - chunk_slots[slot_idx].second = loc; + } else { + __shared__ cuda::atomic counter; + using cuda::std::memory_order_relaxed; + if (t == 0) { new (&counter) cuda::atomic{0}; } + __syncthreads(); + + // Iterate over all slots in the map. + for (; t < chunk.dict_map_size; t += block_size) { + auto* slot = map_storage.data() + chunk.dict_map_offset + t; + auto const key = slot->first; + if (key != KEY_SENTINEL) { + auto const loc = counter.fetch_add(1, memory_order_relaxed); + cudf_assert(loc < MAX_DICT_SIZE && "Number of filled slots exceeds max dict size"); + chunk.dict_data[loc] = key; + // If sorting dict page ever becomes a hard requirement, enable the following statement + // and add a dict sorting step before storing into the slot's second field. + // chunk.dict_data_idx[loc] = idx; + slot->second = loc; + } } } } @@ -424,68 +356,16 @@ CUDF_KERNEL void __launch_bounds__(block_size) s_ck_start_val_idx); } -void populate_chunk_hash_maps(device_span const map_storage, - cudf::detail::device_2dspan frags, - rmm::cuda_stream_view stream) -{ - dim3 const dim_grid(frags.size().second, frags.size().first); - populate_chunk_hash_maps_kernel - <<>>(map_storage, frags); -} - -void collect_map_entries(device_span const map_storage, - device_span chunks, - cudf::detail::device_2dspan frags, - rmm::cuda_stream_view stream) -{ - constexpr int block_size = 1024; - collect_map_entries_kernel - <<>>(map_storage, chunks, frags); -} - -void get_dictionary_indices(device_span const map_storage, - cudf::detail::device_2dspan frags, - rmm::cuda_stream_view stream) -{ - dim3 const dim_grid(frags.size().second, frags.size().first); - get_dictionary_indices_kernel - <<>>(map_storage, frags); -} - -namespace { - -// Warps per block for `compute_page_dict_rle_bits_kernel`. Sized so that each -// block pulls several pages through the reduce, which amortizes the block-level -// overhead over 4 independent warp-level reductions without oversubscribing -// shared memory or hurting occupancy. -constexpr int kDictRleBitsWarpsPerBlock = 4; -constexpr int kDictRleBitsBlockSize = kDictRleBitsWarpsPerBlock * cudf::detail::warp_size; - -// One warp per data page. Each warp strides through its page's slice of -// `chunk->dict_index`, computes the max over *valid* rows only, and stores -// `max(NumRequiredBits(page_max), 1)` into `page.dict_rle_bits`. -// -// Why warp-per-page instead of a `cub::DeviceSegmentedReduce::Max` call per -// chunk (PHASE_2_VARIABLE_BITS.md §2.2.3): segments are small (typically 1k- -// 100k elements per page), so one CUB call per chunk would be launch-overhead -// bound across workloads with many chunks. A single kernel with one warp per -// page keeps the reduction entirely in registers + shuffle and issues exactly -// one launch for the whole writer. -// -// Pages that are skipped (dictionary page itself, non-dict chunks, BOOLEAN -// columns whose dict_bits is always 1) retain the `chunk->dict_rle_bits` -// value that `gpuInitPages` wrote, so the encoder falls back to chunk-wide -// behavior for them. -CUDF_KERNEL void __launch_bounds__(kDictRleBitsBlockSize) +CUDF_KERNEL void __launch_bounds__(DEFAULT_BLOCK_SIZE) compute_page_dict_rle_bits_kernel(device_span pages) { - constexpr auto warp_size = cudf::detail::warp_size; - auto const warp_lane = static_cast(threadIdx.x % warp_size); - auto const warp_id = static_cast(threadIdx.x / warp_size); - auto const page_idx = static_cast(blockIdx.x) * kDictRleBitsWarpsPerBlock + warp_id; + constexpr auto warp_size = cudf::detail::warp_size; + auto const warp_lane = static_cast(threadIdx.x % warp_size); + auto const warp_id = static_cast(threadIdx.x / warp_size); + auto constexpr warps_per_block = DEFAULT_BLOCK_SIZE / warp_size; + auto const page_idx = static_cast(blockIdx.x) * warps_per_block + warp_id; - __shared__ - typename cub::WarpReduce::TempStorage reduce_storage[kDictRleBitsWarpsPerBlock]; + __shared__ typename cub::WarpReduce::TempStorage reduce_storage[warps_per_block]; if (page_idx >= static_cast(pages.size())) { return; } @@ -519,7 +399,7 @@ CUDF_KERNEL void __launch_bounds__(kDictRleBitsBlockSize) for (size_type i = begin + warp_lane; i < end; i += warp_size) { auto const val_idx = chunk_start_val + i; if (val_idx < leaf_size && leaf_col.is_valid(val_idx)) { - lane_max = max(lane_max, dict_index[i]); + lane_max = cuda::std::max(lane_max, dict_index[i]); } } @@ -529,19 +409,51 @@ CUDF_KERNEL void __launch_bounds__(kDictRleBitsBlockSize) if (warp_lane == 0) { // Floor at 1 to match the chunk-wide convention (all-null pages still // emit a 1-bit RLE preamble; see `writer_impl.cu`'s `std::max(..., 1)`). - auto const nbits = max(cuda::std::bit_width(static_cast(page_max)), 1); + auto const nbits = cuda::std::max(cuda::std::bit_width(static_cast(page_max)), 1); page.dict_rle_bits = static_cast(nbits); } } } // namespace +void populate_chunk_hash_maps(device_span const map_storage, + cudf::detail::device_2dspan frags, + rmm::cuda_stream_view stream) +{ + dim3 const dim_grid(frags.size().second, frags.size().first); + populate_chunk_hash_maps_kernel + <<>>(map_storage, frags); +} + +void collect_map_entries(device_span const map_storage, + device_span chunks, + cudf::detail::device_2dspan frags, + rmm::cuda_stream_view stream) +{ + constexpr int block_size = 1024; + static_assert(block_size >= MAX_FRAGMENTS_PER_BLOCK, + "block_size must be >= MAX_FRAGMENTS_PER_BLOCK so one BlockScan thread backs " + "each histogram bucket."); + collect_map_entries_kernel + <<>>(map_storage, chunks, frags); +} + +void get_dictionary_indices(device_span const map_storage, + cudf::detail::device_2dspan frags, + rmm::cuda_stream_view stream) +{ + dim3 const dim_grid(frags.size().second, frags.size().first); + get_dictionary_indices_kernel + <<>>(map_storage, frags); +} + void compute_per_page_dict_rle_bits(device_span pages, rmm::cuda_stream_view stream) { if (pages.empty()) { return; } - auto const num_blocks = cudf::util::div_rounding_up_safe(static_cast(pages.size()), - kDictRleBitsWarpsPerBlock); - compute_page_dict_rle_bits_kernel<<>>( - pages); + auto constexpr warps_per_block = DEFAULT_BLOCK_SIZE / cudf::detail::warp_size; + auto const num_blocks = + cudf::util::div_rounding_up_safe(static_cast(pages.size()), warps_per_block); + compute_page_dict_rle_bits_kernel<<>>(pages); } + } // namespace cudf::io::parquet::detail diff --git a/cpp/src/io/parquet/page_enc.cu b/cpp/src/io/parquet/page_enc.cu index 909cb60f2eee..31f366b39b9d 100644 --- a/cpp/src/io/parquet/page_enc.cu +++ b/cpp/src/io/parquet/page_enc.cu @@ -642,7 +642,6 @@ CUDF_KERNEL void __launch_bounds__(128) uint32_t num_rows = 0; uint32_t page_start = 0; uint32_t page_offset = ck_g.ck_stat_size; - uint32_t num_dict_entries = 0; uint32_t comp_page_offset = ck_g.ck_stat_size; uint32_t page_headers_size = 0; uint32_t max_page_data_size = 0; @@ -897,7 +896,6 @@ CUDF_KERNEL void __launch_bounds__(128) max_stats_len = 0; } max_stats_len = max(max_stats_len, minmax_len); - num_dict_entries += frag_g.num_dict_vals; page_size += fragment_data_size; // fragment_data_size includes the length indicator...remove it var_bytes_size += frag_g.fragment_data_size - frag_g.num_valid * sizeof(size_type); diff --git a/cpp/src/io/parquet/parquet_gpu.cuh b/cpp/src/io/parquet/parquet_gpu.cuh index 3a963adccf31..82ed82617cc3 100644 --- a/cpp/src/io/parquet/parquet_gpu.cuh +++ b/cpp/src/io/parquet/parquet_gpu.cuh @@ -44,9 +44,9 @@ auto constexpr VALUE_SENTINEL = mapped_type{-1}; auto constexpr SCOPE = cuda::thread_scope_block; using storage_type = cuco::bucket_storage, - rmm::mr::polymorphic_allocator>; + bucket_size, + cuco::extent, + rmm::mr::polymorphic_allocator>; using storage_ref_type = typename storage_type::ref_type; /** @@ -100,7 +100,7 @@ inline size_type __device__ row_to_value_idx(size_type idx, * @param stream CUDA stream to use */ void populate_chunk_hash_maps(device_span const map_storage, - cudf::detail::device_2dspan frags, + cudf::detail::device_2dspan frags, rmm::cuda_stream_view stream); /** @@ -156,7 +156,6 @@ void get_dictionary_indices(device_span const map_storage, * @param pages Device span of encoder pages. Field `dict_rle_bits` is written. * @param stream CUDA stream to use */ -void compute_per_page_dict_rle_bits(device_span pages, - rmm::cuda_stream_view stream); +void compute_per_page_dict_rle_bits(device_span pages, rmm::cuda_stream_view stream); } // namespace cudf::io::parquet::detail diff --git a/cpp/src/io/parquet/parquet_gpu.hpp b/cpp/src/io/parquet/parquet_gpu.hpp index 7397a063ba9d..4f7845df1bcf 100644 --- a/cpp/src/io/parquet/parquet_gpu.hpp +++ b/cpp/src/io/parquet/parquet_gpu.hpp @@ -530,7 +530,7 @@ struct PageFragment { uint32_t num_valid; //files[p].key_value_metadata), - [](auto const& kv) { return KeyValue{kv.first, kv.second}; }); + [](auto const& kv) { + return KeyValue{kv.first, kv.second}; + }); } // Append arrow schema to the key-value metadata @@ -1288,7 +1290,7 @@ size_t max_page_bytes(compression_type compression, size_t max_page_size_bytes) std::pair>, std::vector>> build_chunk_dictionaries(hostdevice_2dvector& chunks, host_span col_desc, - device_2dspan frags, + device_2dspan frags, compression_type compression, dictionary_policy dict_policy, size_t max_dict_size, @@ -1885,10 +1887,6 @@ auto convert_table_to_parquet_data(table_input_metadata& table_meta, auto& row_group = agg_meta->file(p).row_groups[global_r]; auto const fragments_in_chunk = util::div_rounding_up_safe(row_group.num_rows, max_page_fragment_size); - CUDF_EXPECTS(fragments_in_chunk <= static_cast(MAX_FRAGMENTS_PER_BLOCK), - "Number of fragments per column chunk exceeds the maximum supported by the " - "parquet writer's dictionary collection kernel. Consider increasing the " - "fragment size or reducing the row group size."); row_group.total_byte_size = 0; row_group.columns.resize(num_columns); for (int c = 0; c < num_columns; c++) { @@ -1902,6 +1900,7 @@ auto convert_table_to_parquet_data(table_input_metadata& table_meta, ck.start_row = start_row; ck.num_rows = (uint32_t)row_group.num_rows; ck.first_fragment = c * num_fragments + f; + ck.num_fragments = fragments_in_chunk; ck.encodings = 0; auto chunk_fragments = row_group_fragments[c].subspan(f, fragments_in_chunk); // In fragment struct, add a pointer to the chunk it belongs to @@ -1966,6 +1965,7 @@ auto convert_table_to_parquet_data(table_input_metadata& table_meta, EncColumnChunk& ck = chunks[r + first_rg_in_part[p]][c]; ck.fragments = page_fragments.device_ptr(frag_offset); ck.first_fragment = frag_offset; + ck.num_fragments = fragments_in_chunk; // update the chunk pointer here for each fragment in chunk.fragments for (uint32_t i = 0; i < fragments_in_chunk; i++) { From 816902b9dc4a13f5f7a99a0bd70db25ebeca7bcc Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 01:16:32 +0000 Subject: [PATCH 09/28] Clean up --- cpp/src/io/parquet/chunk_dict.cu | 9 +++++++++ cpp/src/io/parquet/parquet_gpu.cuh | 9 --------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index 81bbc5cd8747..be684e0591fa 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -24,6 +24,15 @@ namespace cudf::io::parquet::detail { namespace { +// Upper bound on the number of fragments per column chunk that the +// shared-memory histogram in `collect_map_entries_kernel` can accommodate. +// A typical workload is 1M row groups / ~5000-row fragments ≈ 200 fragments +// per chunk, so 1024 is a comfortable ceiling. Host-side code must enforce +// this before launching the kernel (see `build_chunk_dictionaries`); the +// kernel also has a `cudf_assert` as a debug-build safety net, but that is +// compiled out in release builds. +constexpr size_type MAX_FRAGMENTS_PER_BLOCK = 1024; + constexpr int DEFAULT_BLOCK_SIZE = 256; template diff --git a/cpp/src/io/parquet/parquet_gpu.cuh b/cpp/src/io/parquet/parquet_gpu.cuh index 82ed82617cc3..74a9602bf9c2 100644 --- a/cpp/src/io/parquet/parquet_gpu.cuh +++ b/cpp/src/io/parquet/parquet_gpu.cuh @@ -30,15 +30,6 @@ auto constexpr bucket_size = auto constexpr occupancy_factor = 1.43f; ///< cuCollections suggests using a hash map of size ///< N * (1/0.7) = 1.43 to target a 70% occupancy factor. -// Upper bound on the number of fragments per column chunk that the -// shared-memory histogram in `collect_map_entries_kernel` can accommodate. -// A typical workload is 1M row groups / ~5000-row fragments ≈ 200 fragments -// per chunk, so 1024 is a comfortable ceiling. Host-side code must enforce -// this before launching the kernel (see `build_chunk_dictionaries`); the -// kernel also has a `cudf_assert` as a debug-build safety net, but that is -// compiled out in release builds. -constexpr size_type MAX_FRAGMENTS_PER_BLOCK = 1024; - auto constexpr KEY_SENTINEL = key_type{-1}; auto constexpr VALUE_SENTINEL = mapped_type{-1}; auto constexpr SCOPE = cuda::thread_scope_block; From 6686d19fc295278e485ebc7c1c052a1b997351e4 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 16:52:17 +0000 Subject: [PATCH 10/28] humanize gtest --- cpp/tests/io/parquet_writer_test.cpp | 362 ++++++--------------------- 1 file changed, 76 insertions(+), 286 deletions(-) diff --git a/cpp/tests/io/parquet_writer_test.cpp b/cpp/tests/io/parquet_writer_test.cpp index 41a1c0311696..57bfcb552c6b 100644 --- a/cpp/tests/io/parquet_writer_test.cpp +++ b/cpp/tests/io/parquet_writer_test.cpp @@ -1112,300 +1112,90 @@ TEST_F(ParquetWriterTest, SingleValueDictionaryTest) EXPECT_EQ(nbits, expected_bits); } -namespace { - -/** @brief Returns whether a column chunk was actually written with dictionary encoding. - * - * @param chunk The column chunk to check. - * @return Whether the column chunk was actually written with dictionary encoding - */ -[[nodiscard]] bool is_chunk_dict_encoded(cudf::io::parquet::ColumnChunk const& chunk) -{ - return std::any_of( - chunk.meta_data.encodings.begin(), chunk.meta_data.encodings.end(), [](auto const encoding) { - return encoding == cudf::io::parquet::Encoding::PLAIN_DICTIONARY or - encoding == cudf::io::parquet::Encoding::RLE_DICTIONARY; - }); -} - -[[nodiscard]] std::vector read_page_dict_bits( - std::unique_ptr const& source, cudf::io::parquet::ColumnChunk const& chunk) -{ - auto const oi = read_offset_index(source, chunk); - std::vector page_dict_bits; - page_dict_bits.reserve(oi.page_locations.size()); - std::transform( - oi.page_locations.begin(), - oi.page_locations.end(), - std::back_inserter(page_dict_bits), - [&source](auto const& page_location) { return read_dict_bits(source, page_location); }); - return page_dict_bits; -} - -template -[[nodiscard]] auto write_and_read_dict_bits(table_view const& expected, - ConfigureWriter&& configure_writer) - -> std::pair, cudf::io::parquet::FileMetaData> -{ - auto buffer = std::vector{}; - auto builder = cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&buffer}, expected) - .compression(cudf::io::compression_type::NONE) - .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) - .dictionary_policy(cudf::io::dictionary_policy::ALWAYS); - configure_writer(builder); - - cudf::io::parquet_writer_options out_opts = builder; - cudf::io::write_parquet(out_opts); - - auto const buffer_span = - cudf::host_span(reinterpret_cast(buffer.data()), buffer.size()); - cudf::io::parquet_reader_options in_opts = - cudf::io::parquet_reader_options::builder(cudf::io::source_info(buffer_span)); - auto const result = cudf::io::read_parquet(in_opts); - CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); - - auto source = cudf::io::datasource::create(cudf::host_span{buffer_span}); - cudf::io::parquet::FileMetaData fmd; - read_footer(source, &fmd); - - auto const& chunk = fmd.row_groups.front().columns.front(); - return {is_chunk_dict_encoded(chunk) ? read_page_dict_bits(source, chunk) : std::vector{}, - std::move(fmd)}; -} - -} // namespace - -// Phase 2 per-page variable-bit-width RLE coverage (PHASE_2_VARIABLE_BITS.md §2.4). -// The four tests below exercise the new code path from three angles: -// (1) decoder correctness across cardinalities and physical types -// (integers / strings / list), -// (2) the chunk-wide upper bound is respected on every page, -// (3) the optimization actually fires -- file size drops below the pre-Phase-2 -// chunk-wide number on the benchmark workload. - -// Round-trip INT64 dictionary-encoded columns at three cardinalities. The -// cardinalities are chosen to straddle the 10-bit / 16-bit / 20-bit RLE widths -// so the reader has to correctly consume per-page preambles of different -// widths. Chunk-wide-width encoding would pass round-trip as well, so this is -// primarily a "no reader mismatch under variable widths" gate; the upper-bound -// assertion catches an accidental regression that emits *more* bits per page -// than the chunk actually needs. -TEST_F(ParquetWriterTest, VariableBitWidthRoundTripIntegers) -{ - constexpr cudf::size_type nrows = 200'000; - - for (auto const cardinality : {10'000, 64'000, 1'000'000}) { - std::mt19937 rng{0xFEEDFACE}; - std::uniform_int_distribution dist(0, cardinality - 1); - std::vector values(nrows); - std::generate(values.begin(), values.end(), [&] { return dist(rng); }); - - auto const col = cudf::test::fixed_width_column_wrapper(values.begin(), values.end()); - auto const expected = table_view{{col}}; - auto const [page_dict_bits, fmd] = write_and_read_dict_bits(expected, [=](auto& builder) { - builder.row_group_size_rows(nrows).max_page_size_rows(nrows / 4); - }); - - auto const& chunk = fmd.row_groups.front().columns.front(); - // Under `dictionary_policy::ALWAYS` the writer still falls back to PLAIN - // when the dict page would exceed the row data (very sparse cardinality - // relative to row count). Treat that as out-of-scope for bit-width - // checking but still validate round-trip correctness above. - if (not is_chunk_dict_encoded(chunk)) { continue; } - - auto const chunk_wide_max_bits = std::bit_width(cardinality - 1); - ASSERT_GT(page_dict_bits.size(), 1u) << "cardinality=" << cardinality; - for (auto const nbits : page_dict_bits) { - EXPECT_GE(nbits, 1) << "cardinality=" << cardinality; - EXPECT_LE(nbits, chunk_wide_max_bits) << "cardinality=" << cardinality; - } - } -} - -// Round-trip string dictionary encoding at moderate cardinality. Strings take a -// different write path through `build_chunk_dictionaries` (variable-length hash -// keys, dict page is an array of `string_index_pair`), so we cover the encode -// path with a separate fixture rather than folding it into the integer test. -TEST_F(ParquetWriterTest, VariableBitWidthRoundTripStrings) +TEST_F(ParquetWriterTest, VariableBitWidthDictEncoding) { - constexpr cudf::size_type nrows = 100'000; - constexpr cudf::size_type cardinality = 5'000; + constexpr auto num_rows = 100'000; + constexpr auto num_pages = 10; + constexpr auto page_size = num_rows / num_pages; + constexpr auto hot_pages = num_pages - 2; + constexpr auto rare_pages = num_pages - hot_pages; + constexpr auto cardinality = 64'000; + constexpr auto frequent_set_size = 64; - std::mt19937 rng{0xDEADBEEF}; - std::uniform_int_distribution dist(0, cardinality - 1); - std::vector values(nrows); - std::generate(values.begin(), values.end(), [&] { return "str_" + std::to_string(dist(rng)); }); - - auto const col = cudf::test::strings_column_wrapper(values.begin(), values.end()); - auto const expected = table_view{{col}}; - auto const [page_dict_bits, fmd] = write_and_read_dict_bits(expected, [=](auto& builder) { - builder.row_group_size_rows(nrows).max_page_size_rows(nrows / 4); - }); - - auto const chunk_wide_max_bits = - std::max(std::bit_width(static_cast(cardinality - 1)), 1); - ASSERT_TRUE(is_chunk_dict_encoded(fmd.row_groups.front().columns.front())); - ASSERT_GT(page_dict_bits.size(), 1u); - for (auto const nbits : page_dict_bits) { - EXPECT_GE(nbits, 1); - EXPECT_LE(nbits, chunk_wide_max_bits); + auto const filepath = temp_env->get_temp_filepath("VariableBitWidthDictEncoding.parquet"); + { + std::mt19937 rng{0xACAD1A}; + using ColumnType = cudf::test::fixed_width_column_wrapper; + + // Hot pages contain values in [0, frequent_set_size - 1], rare pages contain values in + // [frequent_set_size, cardinality - 1] + std::uniform_int_distribution freq_dist(0, frequent_set_size - 1); + std::uniform_int_distribution rare_dist(frequent_set_size, cardinality - 1); + auto constexpr threshold = hot_pages * page_size; + auto values = cudf::detail::make_counting_transform_iterator( + 0, [&](auto i) { return i < threshold ? freq_dist(rng) : rare_dist(rng); }); + auto const col = ColumnType(values, values + num_rows); + + auto writer_opts = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, table_view{{col}}) + .compression(cudf::io::compression_type::NONE) // No compression + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) // Write page index + .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) // Always dictionary encode + .row_group_size_rows(num_rows) // Single row group only + .max_page_size_rows(page_size) // Max page size is set to the page size + .max_page_size_bytes(std::size_t{64} << 20) // Unlimited page size + .build(); + cudf::io::write_parquet(writer_opts); } -} -// Round-trip list dictionary encoding. Nested types exercise the -// `row_to_value_idx` path inside `compute_page_dict_rle_bits_kernel`: the -// per-page max must be taken over *leaf* values, not rows, and the kernel must -// correctly translate `page.start_row`/`page.num_leaf_values` into leaf-space -// offsets. A mis-translation here would produce a bit width too small to cover -// some leaf dict_index, and the reader would mis-decode the page -- so the -// round-trip comparison below is the real gate. We deliberately don't inspect -// `read_dict_bits` here because the helper assumes a flat column layout (dict -// RLE stream at byte 0 of the page payload); for list the page payload -// starts with rep+def levels and dereferencing byte 0 as a bit width is -// meaningless. -TEST_F(ParquetWriterTest, VariableBitWidthRoundTripLists) -{ - constexpr cudf::size_type num_lists = 20'000; - constexpr cudf::size_type cardinality = 1'024; - - std::mt19937 rng{0xCAFEF00D}; - std::uniform_int_distribution list_len_dist(0, 5); - std::uniform_int_distribution val_dist(0, cardinality - 1); - - std::vector leaf_values; - std::vector offsets{0}; - leaf_values.reserve(num_lists * 3); - offsets.reserve(num_lists + 1); - for (cudf::size_type i = 0; i < num_lists; ++i) { - auto const list_len = list_len_dist(rng); - for (int j = 0; j < list_len; ++j) { - leaf_values.push_back(val_dist(rng)); - } - offsets.push_back(static_cast(leaf_values.size())); - } + auto datasource = cudf::io::datasource::create(filepath); - auto leaf_col = - cudf::test::fixed_width_column_wrapper(leaf_values.begin(), leaf_values.end()) - .release(); - auto offsets_col = - cudf::test::fixed_width_column_wrapper(offsets.begin(), offsets.end()) - .release(); - auto list_col = cudf::make_lists_column( - num_lists, std::move(offsets_col), std::move(leaf_col), 0, rmm::device_buffer{}); - auto const expected = table_view{{*list_col}}; - auto const [page_dict_bits, fmd] = write_and_read_dict_bits(expected, [=](auto& builder) { - builder.row_group_size_rows(num_lists).max_page_size_rows(num_lists / 4); - }); - static_cast(page_dict_bits); - EXPECT_TRUE(is_chunk_dict_encoded(fmd.row_groups.front().columns.front())); -} - -// End-to-end file-size gate. Reproduces the `parquet_write_dict_encoding` -// benchmark workload (8 "common" pages touching only 64 frequent values + 2 -// "rare" pages touching values 64..63999) and asserts: -// (a) common pages bit-pack at their page-local minimum width (`nbits <= -// std::max(std::bit_width(frequent_set_size - 1), 1)`), -// (b) at least one page reaches the chunk-wide width (so dict_rle_bits is -// still the upper bound, not silently clobbered), -// (c) the resulting file is strictly smaller than the Phase 1 baseline -// (2,491,156 bytes at num_rows = 1,000,000). We use the benchmark's exact -// shape but a smaller row count (200K rows, 20K rows/page) so the test -// finishes in well under a second; the proportional savings are the same. -// See PHASE_2_VARIABLE_BITS.md §2.4 for the design rationale. -TEST_F(ParquetWriterTest, VariableBitWidthSmallerThanChunkWide) -{ - constexpr cudf::size_type pages_per_chunk = 10; - constexpr cudf::size_type hot_pages = pages_per_chunk - 2; - constexpr cudf::size_type page_size_rows = 20'000; - constexpr cudf::size_type num_rows = pages_per_chunk * page_size_rows; - constexpr cudf::size_type cardinality = 64'000; - constexpr cudf::size_type frequent_set_size = 64; - - std::mt19937 rng{0xC0DEFACE}; - std::uniform_int_distribution freq_dist(0, frequent_set_size - 1); - std::uniform_int_distribution rare_dist(frequent_set_size, cardinality - 1); - - std::vector values(num_rows); - cudf::size_type const threshold = hot_pages * page_size_rows; - for (cudf::size_type i = 0; i < num_rows; ++i) { - values[i] = i < threshold ? freq_dist(rng) : rare_dist(rng); + // Extract dictionary bit width for each data page from page index + std::vector page_dict_bits; + { + // Read file metadata + cudf::io::parquet::FileMetaData file_metadata; + read_footer(datasource, &file_metadata); + + // Check dictionary encoded pages + auto const& colchunk = file_metadata.row_groups.front().columns.front(); + EXPECT_TRUE(std::any_of(colchunk.meta_data.encodings.begin(), + colchunk.meta_data.encodings.end(), + [](auto const encoding) { + return encoding == cudf::io::parquet::Encoding::PLAIN_DICTIONARY or + encoding == cudf::io::parquet::Encoding::RLE_DICTIONARY; + })); + + auto const oi = read_offset_index(datasource, colchunk); + page_dict_bits.reserve(oi.page_locations.size()); + std::transform(oi.page_locations.begin(), + oi.page_locations.end(), + std::back_inserter(page_dict_bits), + [&datasource](auto const& page_location) { + return read_dict_bits(datasource, page_location); + }); + + // Check page count + EXPECT_EQ(oi.page_locations.size(), static_cast(num_pages)); } - auto const col = cudf::test::fixed_width_column_wrapper(values.begin(), values.end()); - auto const expected = table_view{{col}}; - auto buffer = std::vector{}; - auto const [page_dict_bits, fmd] = [&]() { - auto local_buffer = std::vector{}; - auto builder = - cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&local_buffer}, expected) - .compression(cudf::io::compression_type::NONE) - .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) - .dictionary_policy(cudf::io::dictionary_policy::ALWAYS) - .row_group_size_rows(num_rows) - .max_page_size_rows(page_size_rows) - .max_page_size_bytes(std::size_t{64} << 20); - cudf::io::parquet_writer_options out_opts = builder; - cudf::io::write_parquet(out_opts); - - auto const buffer_span = cudf::host_span( - reinterpret_cast(local_buffer.data()), local_buffer.size()); - cudf::io::parquet_reader_options in_opts = - cudf::io::parquet_reader_options::builder(cudf::io::source_info(buffer_span)); - auto const result = cudf::io::read_parquet(in_opts); - CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view()); - - auto source = cudf::io::datasource::create(cudf::host_span{buffer_span}); - cudf::io::parquet::FileMetaData local_fmd; - read_footer(source, &local_fmd); - - buffer = std::move(local_buffer); - auto const& chunk = local_fmd.row_groups.front().columns.front(); - return std::pair{ - is_chunk_dict_encoded(chunk) ? read_page_dict_bits(source, chunk) : std::vector{}, - std::move(local_fmd)}; - }(); - - EXPECT_TRUE(is_chunk_dict_encoded(fmd.row_groups.front().columns.front())); - - // Chunk-wide upper bound computed from the full cardinality. The writer's - // `dict_rle_bits` is `NumRequiredBits(num_dict_entries - 1)`, and under - // uniform sampling of `[0, cardinality)` the actual number of distinct - // values drawn may land a bit below cardinality -- so we use this value as - // a *ceiling* rather than an equality target. - auto const chunk_wide_max_bits = std::bit_width(cardinality - 1); - auto const frequent_max_bits = std::bit_width(frequent_set_size - 1); // 6 - - ASSERT_EQ(page_dict_bits.size(), static_cast(pages_per_chunk)); - - int common_page_count = 0; - int rare_page_count = 0; - int max_observed_bits = 0; - for (auto const nbits : page_dict_bits) { - EXPECT_GE(nbits, 1); - EXPECT_LE(nbits, chunk_wide_max_bits); - max_observed_bits = std::max(max_observed_bits, nbits); - if (nbits <= frequent_max_bits) { - ++common_page_count; - } else { - ++rare_page_count; - } + // Checks + { + // Check min and max bit widths + auto const [min_bits, max_bits] = std::ranges::minmax_element(page_dict_bits); + auto const chunk_wide_max_bits = std::bit_width(cardinality - 1); + auto const frequent_max_bits = std::bit_width(frequent_set_size - 1); + EXPECT_GT(*min_bits, 1); + EXPECT_GT(*max_bits, frequent_max_bits); + EXPECT_LE(*max_bits, chunk_wide_max_bits); + + // Check expected number of hot and rare pages + auto const total_page_count = static_cast(page_dict_bits.size()); + auto const hot_page_count = static_cast( + std::ranges::count_if(page_dict_bits, [&](int nbits) { return nbits <= frequent_max_bits; })); + EXPECT_EQ(hot_page_count, hot_pages); + EXPECT_EQ(total_page_count - hot_page_count, rare_pages); } - // Under Phase 1 ordering the first `hot_pages` pages reference only the - // frequent-set dict_ids `[0, 64)`, so they should bit-pack to <= 6 bits. - // The 2 rare pages reference dict_ids outside that range and must use - // strictly more bits -- this is the whole optimization in one assertion. - EXPECT_EQ(common_page_count, hot_pages); - EXPECT_EQ(rare_page_count, pages_per_chunk - hot_pages); - EXPECT_GT(max_observed_bits, frequent_max_bits); - - // Pre-Phase-2 (chunk-wide 16 bits on every page) baseline at this workload - // scales linearly from the 1M-row benchmark number (2,491,156 bytes): - // 2,491,156 * (200'000 / 1'000'000) = 498,231 bytes - // With per-page variable widths the 8 common pages drop from 16 to <=6 bits - // and should shave ~250 KB (8 pages * 20k values * (16-6) bits / 8). - // Use a conservative bound of 450,000 bytes to leave headroom for future - // encoder improvements. - EXPECT_LT(buffer.size(), 450'000u); } TEST_F(ParquetWriterTest, DictionaryNeverTest) From 47f7ad3e0455cbcb2a7b2085abeb8c5a843ab991 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 17:05:38 +0000 Subject: [PATCH 11/28] Cleanup --- cpp/src/io/parquet/chunk_dict.cu | 78 ++++++++++++++++---------------- 1 file changed, 40 insertions(+), 38 deletions(-) diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index be684e0591fa..43c5a0e9f419 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -59,23 +59,26 @@ struct hash_functor { template struct map_insert_fn { storage_ref_type const& storage_ref; - EncColumnChunk* const& chunk; - PageFragment* const& frag; + EncColumnChunk* const chunk; + PageFragment* const frag; mapped_type const frag_idx; template - __device__ void operator()(size_type const s_start_value_idx, size_type const end_value_idx) + __device__ void operator()(size_type const start_value_idx, size_type const end_value_idx) { if constexpr (column_device_view::has_element_accessor()) { using block_reduce = cub::BlockReduce; __shared__ typename block_reduce::TempStorage reduce_storage; + namespace cg = cooperative_groups; + + auto const block = cg::this_thread_block(); auto const col = chunk->col_desc; column_device_view const& data_col = *col->leaf_column; __shared__ size_type total_num_dict_entries; __shared__ size_type num_dict_vals; - if (threadIdx.x == 0) { num_dict_vals = 0; } - __syncthreads(); + cg::invoke_one(block, [&]() { num_dict_vals = 0; }); + block.sync(); using equality_fn_type = equality_functor; using hash_fn_type = hash_functor; @@ -100,7 +103,7 @@ struct map_insert_fn { cuda::atomic_ref const chunk_uniq_data_size{chunk->uniq_data_size}; // Note: Adjust the following loop to use `cg::tile` if needed in the future. - for (size_type val_idx = s_start_value_idx + t; val_idx - t < end_value_idx; + for (key_type val_idx = start_value_idx + t; val_idx - t < end_value_idx; val_idx += block_size) { size_type is_unique = 0; size_type uniq_elem_size = 0; @@ -112,7 +115,7 @@ struct map_insert_fn { // Insert tile_val_idx to hash map and count successful insertions. if (is_valid) { // Insert the keys using a single thread for best performance for now. - is_unique = map_insert_ref.insert(slot_type{static_cast(val_idx), frag_idx}); + is_unique = map_insert_ref.insert(slot_type{val_idx, frag_idx}); uniq_elem_size = [&]() -> size_type { if (not is_unique) { return 0; } switch (col->physical_type) { @@ -146,24 +149,24 @@ struct map_insert_fn { } // Reduce num_unique and uniq_data_size from all tiles. auto num_unique = block_reduce(reduce_storage).Sum(is_unique); - __syncthreads(); + block.sync(); auto uniq_data_size = block_reduce(reduce_storage).Sum(uniq_elem_size); - // The first thread in the block atomically updates total num_unique and uniq_data_size, - // and accumulates the per-fragment winning-insert count. - if (t == 0) { + // One thread atomically updates the number and data size of total unique values as well as + // the number of unique values inserted by this fragment. + cg::invoke_one(block, [&]() { total_num_dict_entries = chunk_num_dict_entries.fetch_add(num_unique, cuda::std::memory_order_relaxed); total_num_dict_entries += num_unique; num_dict_vals += num_unique; chunk_uniq_data_size.fetch_add(uniq_data_size, cuda::std::memory_order_relaxed); - } - __syncthreads(); + }); + block.sync(); // Check if the num unique values in chunk has already exceeded max dict size and early exit if (total_num_dict_entries > MAX_DICT_SIZE) { break; } } // for loop - // Flush the number of unique values inserted to the fragment. - if (t == 0) { frag->num_dict_vals = num_dict_vals; } + // Flush the number of unique values inserted by this fragment + cg::invoke_one(block, [&]() { frag->num_dict_vals = num_dict_vals; }); } else { CUDF_UNREACHABLE("Unsupported type to insert in map"); } @@ -173,11 +176,11 @@ struct map_insert_fn { template struct map_find_fn { storage_ref_type const& storage_ref; - EncColumnChunk* const& chunk; + EncColumnChunk* const chunk; template - __device__ void operator()(size_type const s_start_value_idx, + __device__ void operator()(size_type const start_value_idx, size_type const end_value_idx, - size_type const s_ck_start_val_idx) + size_type const ck_start_val_idx) { if constexpr (column_device_view::has_element_accessor()) { auto const col = chunk->col_desc; @@ -202,8 +205,7 @@ struct map_find_fn { auto const t = threadIdx.x; // Note: Adjust the following loop to use `cg::tiles` if needed in the future. - for (thread_index_type val_idx = s_start_value_idx + t; val_idx < end_value_idx; - val_idx += block_size) { + for (key_type val_idx = start_value_idx + t; val_idx < end_value_idx; val_idx += block_size) { // Find the key using a single thread for best performance for now. if (data_col.is_valid(val_idx)) { auto const found_slot = map_find_ref.find(val_idx); @@ -211,7 +213,7 @@ struct map_find_fn { cudf_assert(found_slot != map_find_ref.end() && "Unable to find value in map in dictionary index construction"); // No need for atomic as this is not going to be modified by any other thread. - chunk->dict_index[val_idx - s_ck_start_val_idx] = found_slot->second; + chunk->dict_index[val_idx - ck_start_val_idx] = found_slot->second; } } } else { @@ -228,17 +230,17 @@ CUDF_KERNEL void __launch_bounds__(block_size) auto const col_idx = blockIdx.y; auto const frag_idx = blockIdx.x; auto& frag = frags[col_idx][frag_idx]; - auto chunk = frag.chunk; + auto const chunk = frag.chunk; auto col = chunk->col_desc; if (not chunk->use_dictionary) { return; } - size_type start_row = frag.start_row; - size_type end_row = frag.start_row + frag.num_rows; + auto const start_row = frag.start_row; + auto const end_row = frag.start_row + frag.num_rows; // Find the bounds of values in leaf column to be inserted into the map for current chunk. - size_type const s_start_value_idx = row_to_value_idx(start_row, *col); - size_type const end_value_idx = row_to_value_idx(end_row, *col); + auto const start_value_idx = row_to_value_idx(start_row, *col); + auto const end_value_idx = row_to_value_idx(end_row, *col); column_device_view const& data_col = *col->leaf_column; storage_ref_type const storage_ref{chunk->dict_map_size, @@ -246,7 +248,7 @@ CUDF_KERNEL void __launch_bounds__(block_size) type_dispatcher( data_col.type(), map_insert_fn{storage_ref, chunk, &frag, static_cast(frag_idx)}, - s_start_value_idx, + start_value_idx, end_value_idx); } @@ -338,21 +340,21 @@ CUDF_KERNEL void __launch_bounds__(block_size) get_dictionary_indices_kernel(device_span const map_storage, cudf::detail::device_2dspan frags) { - auto const col_idx = blockIdx.y; - auto const block_x = blockIdx.x; - auto const frag = frags[col_idx][block_x]; - auto chunk = frag.chunk; + auto const col_idx = blockIdx.y; + auto const frag_idx = blockIdx.x; + auto const& frag = frags[col_idx][frag_idx]; + auto const chunk = frag.chunk; if (not chunk->use_dictionary) { return; } - size_type start_row = frag.start_row; - size_type end_row = frag.start_row + frag.num_rows; + auto const start_row = frag.start_row; + auto const end_row = frag.start_row + frag.num_rows; auto const col = chunk->col_desc; // Find the bounds of values in leaf column to be searched in the map for current chunk - auto const s_start_value_idx = row_to_value_idx(start_row, *col); - auto const s_ck_start_val_idx = row_to_value_idx(chunk->start_row, *col); - auto const end_value_idx = row_to_value_idx(end_row, *col); + auto const start_value_idx = row_to_value_idx(start_row, *col); + auto const end_value_idx = row_to_value_idx(end_row, *col); + auto const ck_start_val_idx = row_to_value_idx(chunk->start_row, *col); column_device_view const& data_col = *col->leaf_column; storage_ref_type const storage_ref{chunk->dict_map_size, @@ -360,9 +362,9 @@ CUDF_KERNEL void __launch_bounds__(block_size) type_dispatcher(data_col.type(), map_find_fn{storage_ref, chunk}, - s_start_value_idx, + start_value_idx, end_value_idx, - s_ck_start_val_idx); + ck_start_val_idx); } CUDF_KERNEL void __launch_bounds__(DEFAULT_BLOCK_SIZE) From 47012e208d5f27e28ecf0c9d71c601619e808128 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 17:37:42 +0000 Subject: [PATCH 12/28] Cleanup --- cpp/src/io/parquet/chunk_dict.cu | 44 ++++++++++++++++---------------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index 43c5a0e9f419..1db570c1c4bf 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -70,15 +70,13 @@ struct map_insert_fn { using block_reduce = cub::BlockReduce; __shared__ typename block_reduce::TempStorage reduce_storage; - namespace cg = cooperative_groups; - - auto const block = cg::this_thread_block(); + auto const t = threadIdx.x; auto const col = chunk->col_desc; column_device_view const& data_col = *col->leaf_column; __shared__ size_type total_num_dict_entries; __shared__ size_type num_dict_vals; - cg::invoke_one(block, [&]() { num_dict_vals = 0; }); - block.sync(); + if (t == 0) { num_dict_vals = 0; }; + __syncthreads(); using equality_fn_type = equality_functor; using hash_fn_type = hash_functor; @@ -96,14 +94,13 @@ struct map_insert_fn { // Create a map ref with `cuco::insert` operator auto map_insert_ref = hash_map_ref.rebind_operators(cuco::insert); - auto const t = threadIdx.x; // Create atomic refs to the current chunk's num_dict_entries and uniq_data_size cuda::atomic_ref const chunk_num_dict_entries{chunk->num_dict_entries}; cuda::atomic_ref const chunk_uniq_data_size{chunk->uniq_data_size}; // Note: Adjust the following loop to use `cg::tile` if needed in the future. - for (key_type val_idx = start_value_idx + t; val_idx - t < end_value_idx; + for (size_type val_idx = start_value_idx + t; val_idx - t < end_value_idx; val_idx += block_size) { size_type is_unique = 0; size_type uniq_elem_size = 0; @@ -115,7 +112,7 @@ struct map_insert_fn { // Insert tile_val_idx to hash map and count successful insertions. if (is_valid) { // Insert the keys using a single thread for best performance for now. - is_unique = map_insert_ref.insert(slot_type{val_idx, frag_idx}); + is_unique = map_insert_ref.insert(slot_type{static_cast(val_idx), frag_idx}); uniq_elem_size = [&]() -> size_type { if (not is_unique) { return 0; } switch (col->physical_type) { @@ -149,24 +146,24 @@ struct map_insert_fn { } // Reduce num_unique and uniq_data_size from all tiles. auto num_unique = block_reduce(reduce_storage).Sum(is_unique); - block.sync(); + __syncthreads(); auto uniq_data_size = block_reduce(reduce_storage).Sum(uniq_elem_size); // One thread atomically updates the number and data size of total unique values as well as // the number of unique values inserted by this fragment. - cg::invoke_one(block, [&]() { + if (t == 0) { total_num_dict_entries = chunk_num_dict_entries.fetch_add(num_unique, cuda::std::memory_order_relaxed); total_num_dict_entries += num_unique; num_dict_vals += num_unique; chunk_uniq_data_size.fetch_add(uniq_data_size, cuda::std::memory_order_relaxed); - }); - block.sync(); + } + __syncthreads(); // Check if the num unique values in chunk has already exceeded max dict size and early exit if (total_num_dict_entries > MAX_DICT_SIZE) { break; } } // for loop // Flush the number of unique values inserted by this fragment - cg::invoke_one(block, [&]() { frag->num_dict_vals = num_dict_vals; }); + if (t == 0) { frag->num_dict_vals = num_dict_vals; }; } else { CUDF_UNREACHABLE("Unsupported type to insert in map"); } @@ -291,12 +288,18 @@ CUDF_KERNEL void __launch_bounds__(block_size) using block_scan = cub::BlockScan; __shared__ typename block_scan::TempStorage scan_storage; - auto const per_thread_count = (t < num_frags) ? col_frags[frag_start + t].num_dict_vals : 0; - auto per_thread_offset = 0; - block_scan(scan_storage).ExclusiveSum(per_thread_count, per_thread_offset); - if (t < num_frags) { fragment_cursor[t] = per_thread_offset; } + auto base_idx = 0; + while (base_idx < num_frags) { + auto const idx = base_idx + t; + auto const per_thread_count = + (idx < num_frags) ? col_frags[frag_start + idx].num_dict_vals : 0; + auto per_thread_offset = 0; + block_scan(scan_storage).ExclusiveSum(per_thread_count, per_thread_offset); + if (idx < num_frags) { fragment_cursor[idx] = per_thread_offset; } + base_idx += block_size; + __syncthreads(); + } } - __syncthreads(); // Iterate over all slots in the map, claim a dict_id in the bucket and write it to dict_data for (; t < chunk.dict_map_size; t += block_size) { @@ -326,10 +329,7 @@ CUDF_KERNEL void __launch_bounds__(block_size) auto const loc = counter.fetch_add(1, memory_order_relaxed); cudf_assert(loc < MAX_DICT_SIZE && "Number of filled slots exceeds max dict size"); chunk.dict_data[loc] = key; - // If sorting dict page ever becomes a hard requirement, enable the following statement - // and add a dict sorting step before storing into the slot's second field. - // chunk.dict_data_idx[loc] = idx; - slot->second = loc; + slot->second = loc; } } } From cdc379e3a7d12bd9e261062cffde92a3cc1286c0 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 17:44:45 +0000 Subject: [PATCH 13/28] Cleanup --- cpp/src/io/parquet/page_enc.cu | 21 +++++---------------- cpp/src/io/parquet/parquet_gpu.hpp | 10 ++-------- cpp/src/io/parquet/writer_impl.cu | 8 ++------ 3 files changed, 9 insertions(+), 30 deletions(-) diff --git a/cpp/src/io/parquet/page_enc.cu b/cpp/src/io/parquet/page_enc.cu index 31f366b39b9d..4dae7dac6c98 100644 --- a/cpp/src/io/parquet/page_enc.cu +++ b/cpp/src/io/parquet/page_enc.cu @@ -676,9 +676,7 @@ CUDF_KERNEL void __launch_bounds__(128) page_g.num_rows = ck_g.num_dict_entries; page_g.num_leaf_values = ck_g.num_dict_entries; page_g.num_values = ck_g.num_dict_entries; // TODO: shouldn't matter for dict page - // Unused for the dictionary page itself (guarded in gpuEncodeDictPages), but - // initialized for consistency so the field is never observed zero on a populated page. - page_g.dict_rle_bits = ck_g.dict_rle_bits; + page_g.dict_rle_bits = ck_g.dict_rle_bits; // TODO: shouldn't matter for dict page page_offset += util::round_up_unsafe(page_g.max_hdr_size + page_g.max_data_size, page_align); if (not comp_page_sizes.empty()) { @@ -780,11 +778,9 @@ CUDF_KERNEL void __launch_bounds__(128) page_g.data_size = 0; page_g.comp_data_size = 0; page_g.is_compressed = false; - // Conservative initial value: the chunk-wide bit width. A subsequent pass - // (`compute_per_page_dict_rle_bits`) tightens this to the page-local max - // once page boundaries are fixed and dict_index has been materialized. - page_g.dict_rle_bits = ck_g.dict_rle_bits; - page_g.max_hdr_size = max_data_page_hdr_size; // Max size excluding statistics + page_g.dict_rle_bits = + ck_g.dict_rle_bits; // Conservatively set to the chunk-wide bit width + page_g.max_hdr_size = max_data_page_hdr_size; // Max size excluding statistics if (ck_g.stats) { uint32_t stats_hdr_len = 16; if (col_g.stats_dtype == dtype_string || col_g.stats_dtype == dtype_byte_array) { @@ -1914,16 +1910,9 @@ CUDF_KERNEL void __launch_bounds__(block_size, 8) }(); // TODO assert dict_bits >= 0 - // - // For data pages we use `page.dict_rle_bits`, which `compute_per_page_dict_rle_bits` - // tightens to `ceil(log2(page_max_dict_index + 1))`, not the chunk-wide upper bound - // stored in `chunk.dict_rle_bits`. This is the encode-side half of the per-page - // variable-bit-width optimization (PHASE_2_VARIABLE_BITS.md). Page-size estimation - // in `gpuInitPages` intentionally keeps using `chunk.dict_rle_bits` so buffer sizing - // remains a conservative upper bound -- PROBLEM.md §6. auto const dict_bits = (physical_type == Type::BOOLEAN) ? 1 : (s->ck.use_dictionary and s->page.page_type != PageType::DICTIONARY_PAGE) - ? s->page.dict_rle_bits + ? s->page.dict_rle_bits // Use `page.dict_rle_bits` for data pages : -1; if (t == 0) { uint8_t* dst = s->cur; diff --git a/cpp/src/io/parquet/parquet_gpu.hpp b/cpp/src/io/parquet/parquet_gpu.hpp index 4f7845df1bcf..b594618ca72a 100644 --- a/cpp/src/io/parquet/parquet_gpu.hpp +++ b/cpp/src/io/parquet/parquet_gpu.hpp @@ -578,8 +578,7 @@ struct EncColumnChunk { uint32_t num_rows; //!< Number of rows in chunk size_type num_values; //!< Number of values in chunk. Different from num_rows for nested types uint32_t first_fragment; //!< First fragment of chunk - uint32_t num_fragments; //!< Number of fragments in chunk. Set host-side during row group setup - //!< and consumed by `collect_map_entries_kernel` as the bucket count. + uint32_t num_fragments; //!< Number of fragments in chunk EncPage* pages; //!< Ptr to pages that belong to this chunk uint32_t first_page; //!< First page of chunk uint32_t num_pages; //!< Number of pages in chunk @@ -651,12 +650,7 @@ struct EncPage { Encoding encoding; //!< Encoding used for page data uint16_t num_fragments; //!< Number of fragments in page bool is_compressed; //!< Whether this page is compressed (for V2 page-level compression) - uint8_t dict_rle_bits; //!< RLE bit width for this data page's dict indices. - //!< Initialized to `chunk->dict_rle_bits` (the conservative - //!< chunk-wide bound) and then tightened by - //!< `compute_per_page_dict_rle_bits` to - //!< `max(NumRequiredBits(page_max_dict_index), 1)`. Unused - //!< for the dictionary page itself and for non-dict pages. + uint8_t dict_rle_bits; //!< RLE bit width for this data page's dict indices [[nodiscard]] CUDF_HOST_DEVICE constexpr bool is_v2() const { diff --git a/cpp/src/io/parquet/writer_impl.cu b/cpp/src/io/parquet/writer_impl.cu index 6ab710c30439..bbe6984fe712 100644 --- a/cpp/src/io/parquet/writer_impl.cu +++ b/cpp/src/io/parquet/writer_impl.cu @@ -2123,12 +2123,8 @@ auto convert_table_to_parquet_data(table_input_metadata& table_meta, write_v2_headers, stream); - // Now that page boundaries are finalized and `chunk->dict_index` has been - // materialized by `build_chunk_dictionaries`, tighten each data page's - // `dict_rle_bits` from the chunk-wide conservative bound to the page-local - // minimum required width. Page-size estimation intentionally keeps the - // chunk-wide value (see PROBLEM.md §6): estimation is upstream of this - // pass and must remain conservative so page buffers never overflow. + // Now that page boundaries are finalized and dictionary indices have been materialized, compute + // minimum required RLE bit width for each data page compute_per_page_dict_rle_bits({pages.data(), pages.size()}, stream); } From 589f9d056b03d11b0d8268077b41125e3d8598ec Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 18:55:37 +0000 Subject: [PATCH 14/28] Humanize phase 2 --- cpp/src/io/parquet/chunk_dict.cu | 123 ++++++++++++++++++++--------- cpp/src/io/parquet/parquet_gpu.cuh | 26 ++---- 2 files changed, 91 insertions(+), 58 deletions(-) diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index 1db570c1c4bf..b728d37810e3 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -24,17 +24,22 @@ namespace cudf::io::parquet::detail { namespace { -// Upper bound on the number of fragments per column chunk that the -// shared-memory histogram in `collect_map_entries_kernel` can accommodate. -// A typical workload is 1M row groups / ~5000-row fragments ≈ 200 fragments -// per chunk, so 1024 is a comfortable ceiling. Host-side code must enforce -// this before launching the kernel (see `build_chunk_dictionaries`); the -// kernel also has a `cudf_assert` as a debug-build safety net, but that is -// compiled out in release builds. -constexpr size_type MAX_FRAGMENTS_PER_BLOCK = 1024; +/// Maximum number of chunk fragments for which spatially-local dictionary indices across fragment +/// keys are claimed using shared memory +constexpr size_type MAX_FRAGMENTS_PER_CHUNK = 1024; +/// Default block size for kernel launches constexpr int DEFAULT_BLOCK_SIZE = 256; +/** + * @brief Functor for checking equality of two keys. + * + * @tparam T Type of the keys + * @param col Column view + * @param lhs_idx Index of the left-hand side key + * @param rhs_idx Index of the right-hand side key + * @return Whether the keys are equal + */ template struct equality_functor { column_device_view const& col; @@ -46,6 +51,14 @@ struct equality_functor { } }; +/** + * @brief Functor for hashing a key + * + * @tparam T Type of the key + * @param col Column view + * @param idx Key index of the key + * @return Hash value of the key + */ template struct hash_functor { column_device_view const& col; @@ -56,6 +69,15 @@ struct hash_functor { } }; +/** + * @brief Functor for inserting a key into a hash map + * + * @tparam block_size Thread block size + * @param storage_ref Hash map storage reference + * @param chunk Column chunk pointer + * @param frag Page fragment pointer + * @param frag_idx Page fragment index + */ template struct map_insert_fn { storage_ref_type const& storage_ref; @@ -170,6 +192,13 @@ struct map_insert_fn { } }; +/** + * @brief Functor for finding a key in a hash map + * + * @tparam block_size Thread block size + * @param storage_ref Hash map storage reference + * @param chunk Column chunk pointerk + */ template struct map_find_fn { storage_ref_type const& storage_ref; @@ -219,6 +248,13 @@ struct map_find_fn { } }; +/** + * @brief Populate the hash maps for all chunks (thread block per page fragment) + * + * @tparam block_size Thread block size + * @param map_storage Hash map storage span + * @param frags 2D span of page fragments + */ template CUDF_KERNEL void __launch_bounds__(block_size) populate_chunk_hash_maps_kernel(device_span const map_storage, @@ -249,6 +285,14 @@ CUDF_KERNEL void __launch_bounds__(block_size) end_value_idx); } +/** + * @brief Collect the dictionary indices for all chunks (thread block per column chunk) + * + * @tparam block_size Thread block size + * @param map_storage Hash map storage span + * @param chunks Column chunks span + * @param frags 2D span of page fragments + */ template CUDF_KERNEL void __launch_bounds__(block_size) collect_map_entries_kernel(device_span const map_storage, @@ -258,32 +302,19 @@ CUDF_KERNEL void __launch_bounds__(block_size) auto& chunk = chunks[blockIdx.x]; if (not chunk.use_dictionary) { return; } - auto t = threadIdx.x; - - // Resolve the chunk's column-relative fragment range [frag_start, frag_start + num_frags). - // Both values come directly from host-populated fields on the chunk; no in-kernel reduction - // or shared memory is needed. `chunk.fragments` points into `col_frags` (set in - // writer_impl.cu during row_group_fragments setup, before build_chunk_dictionaries runs), - // so the subtraction is the first fragment index stamped into any of this chunk's slots - // by populate_chunk_hash_maps_kernel. `chunk.num_fragments` is the run length. + auto t = threadIdx.x; auto const num_frags = chunk.num_fragments; - if (num_frags <= MAX_FRAGMENTS_PER_BLOCK) { + // If a chunk has less fragments than the maximum fragments per chunk, resolve its + // column-relative fragment range. + if (num_frags <= MAX_FRAGMENTS_PER_CHUNK) { auto const col_idx = chunk.col_desc_id; auto const& col_frags = frags[col_idx]; auto const frag_start = static_cast(chunk.fragments - col_frags.data()); - // Per-bucket cursors: initialized to the exclusive-prefix offsets of the - // per-fragment winning-insert counts; threads `atomicAdd` into these to - // claim dict_ids in Pass 2. - __shared__ size_type fragment_cursor[MAX_FRAGMENTS_PER_BLOCK]; - - // Pass 1: in-block exclusive scan over the per-fragment winning-insert - // counts populate wrote into each fragment. This replaces the old - // slot-rescan histogram: `frag.num_dict_vals` already equals - // "number of slots in `chunk_slots[]` stamped with column-relative fragment - // index (f_start + i)", because populate's `blockIdx.x` is exactly that - // fragment index and each winning CAS contributes exactly one stamp. + // Initialize fragment_offsets with prefix sum (exclusive) of number of dictionary values + // inserted by each fragment in this chunk. + __shared__ size_type fragment_offsets[MAX_FRAGMENTS_PER_CHUNK]; { using block_scan = cub::BlockScan; __shared__ typename block_scan::TempStorage scan_storage; @@ -295,13 +326,14 @@ CUDF_KERNEL void __launch_bounds__(block_size) (idx < num_frags) ? col_frags[frag_start + idx].num_dict_vals : 0; auto per_thread_offset = 0; block_scan(scan_storage).ExclusiveSum(per_thread_count, per_thread_offset); - if (idx < num_frags) { fragment_cursor[idx] = per_thread_offset; } + if (idx < num_frags) { fragment_offsets[idx] = per_thread_offset; } base_idx += block_size; __syncthreads(); } } - // Iterate over all slots in the map, claim a dict_id in the bucket and write it to dict_data + // Iterate over slots and claim a dictionary index from the fragment offsets (spatial-locality + // for keys in the same fragment). for (; t < chunk.dict_map_size; t += block_size) { auto* slot = map_storage.data() + chunk.dict_map_offset + t; auto const key = slot->first; @@ -309,19 +341,21 @@ CUDF_KERNEL void __launch_bounds__(block_size) auto const frag_loc = static_cast(slot->second) - frag_start; cudf_assert(frag_loc >= 0 && frag_loc < num_frags && "populate stamped a fragment hint outside this chunk's fragment range"); - auto const loc = atomicAdd(&fragment_cursor[frag_loc], 1); + auto const loc = atomicAdd(&fragment_offsets[frag_loc], 1); cudf_assert(loc < MAX_DICT_SIZE && "Number of filled slots exceeds max dict size"); chunk.dict_data[loc] = key; slot->second = loc; } } - } else { + } + // Otherwise, iterate over slots and claim monotonically increasing dictionary indices for each + // key + else { __shared__ cuda::atomic counter; using cuda::std::memory_order_relaxed; if (t == 0) { new (&counter) cuda::atomic{0}; } __syncthreads(); - // Iterate over all slots in the map. for (; t < chunk.dict_map_size; t += block_size) { auto* slot = map_storage.data() + chunk.dict_map_offset + t; auto const key = slot->first; @@ -335,6 +369,13 @@ CUDF_KERNEL void __launch_bounds__(block_size) } } +/** + * @brief Get the dictionary indices for all chunks (thread block per page fragment) + * + * @tparam block_size Thread block size + * @param map_storage Hash map storage span + * @param frags 2D span of page fragments + */ template CUDF_KERNEL void __launch_bounds__(block_size) get_dictionary_indices_kernel(device_span const map_storage, @@ -367,8 +408,14 @@ CUDF_KERNEL void __launch_bounds__(block_size) ck_start_val_idx); } +/** + * @brief Compute the RLE bits for the dictionary indices for all pages (warp per page) + * + * @tparam block_size Thread block size + * @param pages Pages span + */ CUDF_KERNEL void __launch_bounds__(DEFAULT_BLOCK_SIZE) - compute_page_dict_rle_bits_kernel(device_span pages) + compute_page_dict_bits_kernel(device_span pages) { constexpr auto warp_size = cudf::detail::warp_size; auto const warp_lane = static_cast(threadIdx.x % warp_size); @@ -442,8 +489,8 @@ void collect_map_entries(device_span const map_storage, rmm::cuda_stream_view stream) { constexpr int block_size = 1024; - static_assert(block_size >= MAX_FRAGMENTS_PER_BLOCK, - "block_size must be >= MAX_FRAGMENTS_PER_BLOCK so one BlockScan thread backs " + static_assert(block_size >= MAX_FRAGMENTS_PER_CHUNK, + "block_size must be >= MAX_FRAGMENTS_PER_CHUNK so one BlockScan thread backs " "each histogram bucket."); collect_map_entries_kernel <<>>(map_storage, chunks, frags); @@ -458,13 +505,13 @@ void get_dictionary_indices(device_span const map_storage, <<>>(map_storage, frags); } -void compute_per_page_dict_rle_bits(device_span pages, rmm::cuda_stream_view stream) +void compute_per_page_dict_bits(device_span pages, rmm::cuda_stream_view stream) { if (pages.empty()) { return; } auto constexpr warps_per_block = DEFAULT_BLOCK_SIZE / cudf::detail::warp_size; auto const num_blocks = cudf::util::div_rounding_up_safe(static_cast(pages.size()), warps_per_block); - compute_page_dict_rle_bits_kernel<<>>(pages); + compute_page_dict_bits_kernel<<>>(pages); } } // namespace cudf::io::parquet::detail diff --git a/cpp/src/io/parquet/parquet_gpu.cuh b/cpp/src/io/parquet/parquet_gpu.cuh index 74a9602bf9c2..d7ddb93a5849 100644 --- a/cpp/src/io/parquet/parquet_gpu.cuh +++ b/cpp/src/io/parquet/parquet_gpu.cuh @@ -97,18 +97,12 @@ void populate_chunk_hash_maps(device_span const map_storage, /** * @brief Compact dictionary hash map entries into chunk.dict_data * - * Each chunk's `dict_id`s are assigned monotonically in fragment-first- - * appearance order: values whose slot was first won by an earlier fragment - * get strictly lower `dict_id`s than values first won by a later fragment. - * This is the ordering prerequisite that lets per-page RLE bit widths shrink - * for pages that only touch values in early fragments. + * `dict_id`s are assigned in fragment-first-insert order, so pages over earlier + * fragments see smaller ids and can use narrower RLE bit widths. * * @param map_storage Bulk hashmap storage * @param chunks Flat span of chunks to compact hash maps for - * @param frags 2D span of per-column page fragments. Used by the collect - * kernel to determine each chunk's column-relative fragment - * range; the span itself matches the one passed to - * `populate_chunk_hash_maps`. + * @param frags 2D span of per-column page fragments * @param stream CUDA stream to use */ void collect_map_entries(device_span const map_storage, @@ -134,19 +128,11 @@ void get_dictionary_indices(device_span const map_storage, rmm::cuda_stream_view stream); /** - * @brief Tighten each data page's `dict_rle_bits` to the minimum width required - * by the max `dict_index` observed in that page's rows. + * @brief Compute the minimum width required for the dictionary indices for each data page * - * Must be invoked after `InitEncoderPages` has finalized page boundaries *and* - * `get_dictionary_indices` has materialized per-chunk `dict_index` arrays, but - * before `EncodePages` reads `page.dict_rle_bits`. Pages that do not use - * dictionary encoding (non-dict chunks, the dictionary page itself, BOOLEAN - * columns) are skipped and keep the `chunk->dict_rle_bits` fallback that - * `gpuInitPages` wrote during page initialization. - * - * @param pages Device span of encoder pages. Field `dict_rle_bits` is written. + * @param pages Device span of encoder pages * @param stream CUDA stream to use */ -void compute_per_page_dict_rle_bits(device_span pages, rmm::cuda_stream_view stream); +void compute_per_page_dict_bits(device_span pages, rmm::cuda_stream_view stream); } // namespace cudf::io::parquet::detail From 01c41f684b312a0f02d1f0d0afc8964c703c88eb Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 18:57:24 +0000 Subject: [PATCH 15/28] Style fix --- cpp/benchmarks/CMakeLists.txt | 4 +--- cpp/src/io/parquet/chunk_dict.cu | 2 +- cpp/src/io/parquet/parquet_gpu.cuh | 8 ++++---- cpp/src/io/parquet/writer_impl.cu | 4 +--- 4 files changed, 7 insertions(+), 11 deletions(-) diff --git a/cpp/benchmarks/CMakeLists.txt b/cpp/benchmarks/CMakeLists.txt index 585266c50a05..e8bc5ae97a00 100644 --- a/cpp/benchmarks/CMakeLists.txt +++ b/cpp/benchmarks/CMakeLists.txt @@ -284,9 +284,7 @@ ConfigureNVBench(BITMASK_NVBENCH bitmask/bitmask_and.cpp bitmask/set_null_mask.c # ################################################################################################## # * parquet writer benchmark ---------------------------------------------------------------------- ConfigureNVBench( - PARQUET_WRITER_NVBENCH - io/parquet/parquet_writer.cpp - io/parquet/parquet_writer_chunks.cpp + PARQUET_WRITER_NVBENCH io/parquet/parquet_writer.cpp io/parquet/parquet_writer_chunks.cpp io/parquet/parquet_writer_dict.cpp ) diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index b728d37810e3..319e7516ba0f 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/io/parquet/parquet_gpu.cuh b/cpp/src/io/parquet/parquet_gpu.cuh index d7ddb93a5849..a57e405ec6c0 100644 --- a/cpp/src/io/parquet/parquet_gpu.cuh +++ b/cpp/src/io/parquet/parquet_gpu.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -35,9 +35,9 @@ auto constexpr VALUE_SENTINEL = mapped_type{-1}; auto constexpr SCOPE = cuda::thread_scope_block; using storage_type = cuco::bucket_storage, - rmm::mr::polymorphic_allocator>; + bucket_size, + cuco::extent, + rmm::mr::polymorphic_allocator>; using storage_ref_type = typename storage_type::ref_type; /** diff --git a/cpp/src/io/parquet/writer_impl.cu b/cpp/src/io/parquet/writer_impl.cu index 369efa6949b1..5bde1a6dc1b8 100644 --- a/cpp/src/io/parquet/writer_impl.cu +++ b/cpp/src/io/parquet/writer_impl.cu @@ -97,9 +97,7 @@ struct aggregate_writer_metadata { std::transform(kv_md[p].begin(), kv_md[p].end(), std::back_inserter(this->files[p].key_value_metadata), - [](auto const& kv) { - return KeyValue{kv.first, kv.second}; - }); + [](auto const& kv) { return KeyValue{kv.first, kv.second}; }); } // Append arrow schema to the key-value metadata From 46761ec07e0a21dade2fbed50856d7f943a9cb45 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 19:09:42 +0000 Subject: [PATCH 16/28] Humanize phase 3 --- cpp/src/io/parquet/chunk_dict.cu | 59 ++++++++++++++++---------------- 1 file changed, 29 insertions(+), 30 deletions(-) diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index 319e7516ba0f..f0cf19301556 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -8,13 +8,15 @@ #include #include #include +#include #include #include #include +#include +#include #include -#include #include #include #include @@ -417,27 +419,25 @@ CUDF_KERNEL void __launch_bounds__(block_size) CUDF_KERNEL void __launch_bounds__(DEFAULT_BLOCK_SIZE) compute_page_dict_bits_kernel(device_span pages) { - constexpr auto warp_size = cudf::detail::warp_size; - auto const warp_lane = static_cast(threadIdx.x % warp_size); - auto const warp_id = static_cast(threadIdx.x / warp_size); - auto constexpr warps_per_block = DEFAULT_BLOCK_SIZE / warp_size; - auto const page_idx = static_cast(blockIdx.x) * warps_per_block + warp_id; - - __shared__ typename cub::WarpReduce::TempStorage reduce_storage[warps_per_block]; + namespace cg = cooperative_groups; + using cudf::detail::warp_size; + auto const page_idx = cudf::detail::grid_1d::global_thread_id() / warp_size; if (page_idx >= static_cast(pages.size())) { return; } - auto& page = pages[page_idx]; - auto const* chunk = page.chunk; - // Non-dict chunk: `dict_rle_bits` is unused by the encoder; skip. + auto const warp = cg::tiled_partition(cg::this_thread_block()); + auto& page = pages[page_idx]; + auto const chunk = page.chunk; + + // Return if non-dict chunk: `dict_rle_bits` is unused by the encoder if (not chunk->use_dictionary) { return; } - // Dictionary page itself does not encode dict_indices; skip. + // Return if dictionary page itself does not encode dict_indices if (page.page_type == PageType::DICTIONARY_PAGE) { return; } - auto const* col = chunk->col_desc; - // BOOLEAN columns emit dict_bits=1 through a separate code path in - // `gpuEncodeDictPages`, independent of `dict_rle_bits`. Leave the field at - // its chunk-wide init value (it is ignored for booleans). + auto const col = chunk->col_desc; + + // Return if BOOLEAN column: dict_bits=1 through a separate code path in + // `gpuEncodeDictPages`, independent of `dict_rle_bits` if (col->physical_type == Type::BOOLEAN) { return; } auto const chunk_start_val = row_to_value_idx(chunk->start_row, *col); @@ -446,30 +446,29 @@ CUDF_KERNEL void __launch_bounds__(DEFAULT_BLOCK_SIZE) auto const begin = page_start_val - chunk_start_val; auto const end = begin + page_num_leaf_values; - auto const* dict_index = chunk->dict_index; - column_device_view const& leaf_col = *col->leaf_column; - auto const leaf_size = leaf_col.size(); + auto const* dict_index = chunk->dict_index; + auto const& leaf_col = *col->leaf_column; + auto const leaf_size = leaf_col.size(); - // Accumulate per-lane max. Null rows leave `dict_index` undefined; gate - // the read with the column's validity bitmap to avoid pulling garbage - // bits into the max. + // Accumulate per-lane max dict index for this page size_type lane_max = 0; - for (size_type i = begin + warp_lane; i < end; i += warp_size) { + for (size_type i = begin + warp.thread_rank(); i < end; i += warp_size) { auto const val_idx = chunk_start_val + i; + // Null rows leave `dict_index` undefined; gate the read with the column's validity bitmap to + // avoid pulling garbage bits into the max. if (val_idx < leaf_size && leaf_col.is_valid(val_idx)) { lane_max = cuda::std::max(lane_max, dict_index[i]); } } - auto const page_max = - cub::WarpReduce(reduce_storage[warp_id]).Reduce(lane_max, cuda::maximum{}); - - if (warp_lane == 0) { - // Floor at 1 to match the chunk-wide convention (all-null pages still - // emit a 1-bit RLE preamble; see `writer_impl.cu`'s `std::max(..., 1)`). + // Write this page's RLE bits + auto const page_max = cg::reduce(warp, lane_max, cg::greater{}); + cg::invoke_one(warp, [&] { + // Floor at 1 to match the chunk-wide convention (all-null pages still emit a 1-bit RLE + // preamble) auto const nbits = cuda::std::max(cuda::std::bit_width(static_cast(page_max)), 1); page.dict_rle_bits = static_cast(nbits); - } + }); } } // namespace From 3ffbf46d255d1dcbb7dfd87f59bbc80ad7fe78ed Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 19:23:54 +0000 Subject: [PATCH 17/28] Add todo --- cpp/src/io/parquet/chunk_dict.cu | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index f0cf19301556..3e4cb0c42cea 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -133,9 +133,12 @@ struct map_insert_fn { auto const is_valid = val_idx < end_value_idx and val_idx < data_col.size() and data_col.is_valid(val_idx); - // Insert tile_val_idx to hash map and count successful insertions. + // Insert fragment index to hash map using a single thread (for best performance for now) + // and count successful insertions. if (is_valid) { - // Insert the keys using a single thread for best performance for now. + // TODO(mh): Here we insert the fragment index of the CAS winner, which may not be the + // smallest one (relies on monotonic block scheduling). Switch to static_map's + // `insert_or_apply` with `cuco::op::min` for deterministic first-fragment semantics is_unique = map_insert_ref.insert(slot_type{static_cast(val_idx), frag_idx}); uniq_elem_size = [&]() -> size_type { if (not is_unique) { return 0; } From 818d9a08abdca75b7539a80ad5f71ec2b862fd09 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 19:24:10 +0000 Subject: [PATCH 18/28] Minor --- cpp/src/io/parquet/writer_impl.cu | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/cpp/src/io/parquet/writer_impl.cu b/cpp/src/io/parquet/writer_impl.cu index 5bde1a6dc1b8..5dbfa4127a24 100644 --- a/cpp/src/io/parquet/writer_impl.cu +++ b/cpp/src/io/parquet/writer_impl.cu @@ -97,7 +97,9 @@ struct aggregate_writer_metadata { std::transform(kv_md[p].begin(), kv_md[p].end(), std::back_inserter(this->files[p].key_value_metadata), - [](auto const& kv) { return KeyValue{kv.first, kv.second}; }); + [](auto const& kv) { + return KeyValue{kv.first, kv.second}; + }); } // Append arrow schema to the key-value metadata @@ -2124,7 +2126,7 @@ auto convert_table_to_parquet_data(table_input_metadata& table_meta, // Now that page boundaries are finalized and dictionary indices have been materialized, compute // minimum required RLE bit width for each data page - compute_per_page_dict_rle_bits({pages.data(), pages.size()}, stream); + compute_per_page_dict_bits({pages.data(), pages.size()}, stream); } // Check device write support for all chunks and initialize bounce_buffer. From e84dfc612986799283128dab92427d392a39f34c Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 19:26:28 +0000 Subject: [PATCH 19/28] Style again --- cpp/src/io/parquet/writer_impl.cu | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/cpp/src/io/parquet/writer_impl.cu b/cpp/src/io/parquet/writer_impl.cu index 5dbfa4127a24..6d0868e72bce 100644 --- a/cpp/src/io/parquet/writer_impl.cu +++ b/cpp/src/io/parquet/writer_impl.cu @@ -97,9 +97,7 @@ struct aggregate_writer_metadata { std::transform(kv_md[p].begin(), kv_md[p].end(), std::back_inserter(this->files[p].key_value_metadata), - [](auto const& kv) { - return KeyValue{kv.first, kv.second}; - }); + [](auto const& kv) { return KeyValue{kv.first, kv.second}; }); } // Append arrow schema to the key-value metadata From 29548709ac708445d855474233121b93d710c0ad Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 19:47:46 +0000 Subject: [PATCH 20/28] Copilot's comments --- cpp/benchmarks/io/parquet/parquet_writer_dict.cpp | 1 + cpp/src/io/parquet/chunk_dict.cu | 2 +- cpp/tests/io/parquet_writer_test.cpp | 11 ++++++++--- 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp b/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp index bbb967a4d1bf..4b61809d17af 100644 --- a/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp +++ b/cpp/benchmarks/io/parquet/parquet_writer_dict.cpp @@ -24,6 +24,7 @@ #include #include +#include #include #include #include diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index 3e4cb0c42cea..7475dd86851b 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -465,7 +465,7 @@ CUDF_KERNEL void __launch_bounds__(DEFAULT_BLOCK_SIZE) } // Write this page's RLE bits - auto const page_max = cg::reduce(warp, lane_max, cg::greater{}); + auto const page_max = cg::reduce(warp, lane_max, cuda::std::greater{}); cg::invoke_one(warp, [&] { // Floor at 1 to match the chunk-wide convention (all-null pages still emit a 1-bit RLE // preamble) diff --git a/cpp/tests/io/parquet_writer_test.cpp b/cpp/tests/io/parquet_writer_test.cpp index 892612394b55..a409167fe80d 100644 --- a/cpp/tests/io/parquet_writer_test.cpp +++ b/cpp/tests/io/parquet_writer_test.cpp @@ -1132,9 +1132,14 @@ TEST_F(ParquetWriterTest, VariableBitWidthDictEncoding) std::uniform_int_distribution freq_dist(0, frequent_set_size - 1); std::uniform_int_distribution rare_dist(frequent_set_size, cardinality - 1); auto constexpr threshold = hot_pages * page_size; - auto values = cudf::detail::make_counting_transform_iterator( - 0, [&](auto i) { return i < threshold ? freq_dist(rng) : rare_dist(rng); }); - auto const col = ColumnType(values, values + num_rows); + auto values = std::vector{}; + values.reserve(num_rows); + std::transform( + cuda::counting_iterator(0), + cuda::counting_iterator(num_rows), + std::back_inserter(values), + [&](auto row_idx) { return row_idx < threshold ? freq_dist(rng) : rare_dist(rng); }); + auto const col = ColumnType(values.begin(), values.end()); auto writer_opts = cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, table_view{{col}}) From db808a7bbe8d004d967285d7bb55576a47c621bb Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 19:49:23 +0000 Subject: [PATCH 21/28] Minor --- cpp/src/io/parquet/chunk_dict.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index 7475dd86851b..87d92acc07fa 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -465,7 +465,7 @@ CUDF_KERNEL void __launch_bounds__(DEFAULT_BLOCK_SIZE) } // Write this page's RLE bits - auto const page_max = cg::reduce(warp, lane_max, cuda::std::greater{}); + auto const page_max = cg::reduce(warp, lane_max, cuda::std::greater{}); cg::invoke_one(warp, [&] { // Floor at 1 to match the chunk-wide convention (all-null pages still emit a 1-bit RLE // preamble) From b9209e969c0d01d14df0ab03b8651e36e0bdd8d7 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 19:52:11 +0000 Subject: [PATCH 22/28] Minor --- cpp/src/io/parquet/chunk_dict.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index 87d92acc07fa..e7800da5e5ca 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -465,7 +465,7 @@ CUDF_KERNEL void __launch_bounds__(DEFAULT_BLOCK_SIZE) } // Write this page's RLE bits - auto const page_max = cg::reduce(warp, lane_max, cuda::std::greater{}); + auto const page_max = cg::reduce(warp, lane_max, cg::greater{}); cg::invoke_one(warp, [&] { // Floor at 1 to match the chunk-wide convention (all-null pages still emit a 1-bit RLE // preamble) From ccc7bf99f097370c0bfe7df758b30108ae18fd6c Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb <14217455+mhaseeb123@users.noreply.github.com> Date: Tue, 28 Apr 2026 13:52:18 -0700 Subject: [PATCH 23/28] Apply suggestion from @mhaseeb123 --- cpp/tests/io/parquet_misc_test.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tests/io/parquet_misc_test.cpp b/cpp/tests/io/parquet_misc_test.cpp index 67803f1592d4..f37cb8ccfc6e 100644 --- a/cpp/tests/io/parquet_misc_test.cpp +++ b/cpp/tests/io/parquet_misc_test.cpp @@ -182,7 +182,7 @@ TEST_P(ParquetSizedTest, DictionaryTest) [&source](auto const& page_location) { return read_dict_bits(source, page_location); }); auto const max_bits = std::max_element(page_dict_bits.begin(), page_dict_bits.end()); - EXPECT_NE(max_bits, page_dict_bits.end()); + ASSERT_NE(max_bits, page_dict_bits.end()); EXPECT_EQ(*max_bits, GetParam()); } From 9a755c173906d13dbb5a25d63784783a15f5333e Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Tue, 28 Apr 2026 20:54:19 +0000 Subject: [PATCH 24/28] Use freq instead of hot to match with rare --- cpp/tests/io/parquet_writer_test.cpp | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/cpp/tests/io/parquet_writer_test.cpp b/cpp/tests/io/parquet_writer_test.cpp index a409167fe80d..7576ca13dd4b 100644 --- a/cpp/tests/io/parquet_writer_test.cpp +++ b/cpp/tests/io/parquet_writer_test.cpp @@ -1117,8 +1117,8 @@ TEST_F(ParquetWriterTest, VariableBitWidthDictEncoding) constexpr auto num_rows = 100'000; constexpr auto num_pages = 10; constexpr auto page_size = num_rows / num_pages; - constexpr auto hot_pages = num_pages - 2; - constexpr auto rare_pages = num_pages - hot_pages; + constexpr auto freq_pages = num_pages - 2; + constexpr auto rare_pages = num_pages - freq_pages; constexpr auto cardinality = 64'000; constexpr auto frequent_set_size = 64; @@ -1131,7 +1131,7 @@ TEST_F(ParquetWriterTest, VariableBitWidthDictEncoding) // [frequent_set_size, cardinality - 1] std::uniform_int_distribution freq_dist(0, frequent_set_size - 1); std::uniform_int_distribution rare_dist(frequent_set_size, cardinality - 1); - auto constexpr threshold = hot_pages * page_size; + auto constexpr threshold = freq_pages * page_size; auto values = std::vector{}; values.reserve(num_rows); std::transform( @@ -1190,16 +1190,18 @@ TEST_F(ParquetWriterTest, VariableBitWidthDictEncoding) auto const [min_bits, max_bits] = std::ranges::minmax_element(page_dict_bits); auto const chunk_wide_max_bits = std::bit_width(cardinality - 1); auto const frequent_max_bits = std::bit_width(frequent_set_size - 1); + ASSERT_NE(min_bits, page_dict_bits.end()); + ASSERT_NE(max_bits, page_dict_bits.end()); EXPECT_GT(*min_bits, 1); EXPECT_GT(*max_bits, frequent_max_bits); EXPECT_LE(*max_bits, chunk_wide_max_bits); - // Check expected number of hot and rare pages + // Check expected number of freq and rare pages auto const total_page_count = static_cast(page_dict_bits.size()); - auto const hot_page_count = static_cast( + auto const freq_page_count = static_cast( std::ranges::count_if(page_dict_bits, [&](int nbits) { return nbits <= frequent_max_bits; })); - EXPECT_EQ(hot_page_count, hot_pages); - EXPECT_EQ(total_page_count - hot_page_count, rare_pages); + EXPECT_EQ(freq_page_count, freq_pages); + EXPECT_EQ(total_page_count - freq_page_count, rare_pages); } } From 2aa81ed6481a9750a6999519ed9f71cec18156f8 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Wed, 29 Apr 2026 00:03:12 +0000 Subject: [PATCH 25/28] Deterministic approach --- cpp/src/io/parquet/chunk_dict.cu | 71 ++++++++++++++++-------------- cpp/src/io/parquet/page_enc.cu | 1 - cpp/src/io/parquet/parquet_gpu.hpp | 1 - 3 files changed, 39 insertions(+), 34 deletions(-) diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index e7800da5e5ca..651dad5b2f32 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -18,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -98,9 +99,6 @@ struct map_insert_fn { auto const col = chunk->col_desc; column_device_view const& data_col = *col->leaf_column; __shared__ size_type total_num_dict_entries; - __shared__ size_type num_dict_vals; - if (t == 0) { num_dict_vals = 0; }; - __syncthreads(); using equality_fn_type = equality_functor; using hash_fn_type = hash_functor; @@ -116,8 +114,8 @@ struct map_insert_fn { cuco::thread_scope_block, storage_ref}; - // Create a map ref with `cuco::insert` operator - auto map_insert_ref = hash_map_ref.rebind_operators(cuco::insert); + // Create a map ref with `cuco::insert_or_apply` operator + auto map_insert_ref = hash_map_ref.rebind_operators(cuco::insert_or_apply); // Create atomic refs to the current chunk's num_dict_entries and uniq_data_size cuda::atomic_ref const chunk_num_dict_entries{chunk->num_dict_entries}; @@ -136,10 +134,9 @@ struct map_insert_fn { // Insert fragment index to hash map using a single thread (for best performance for now) // and count successful insertions. if (is_valid) { - // TODO(mh): Here we insert the fragment index of the CAS winner, which may not be the - // smallest one (relies on monotonic block scheduling). Switch to static_map's - // `insert_or_apply` with `cuco::op::min` for deterministic first-fragment semantics - is_unique = map_insert_ref.insert(slot_type{static_cast(val_idx), frag_idx}); + // Insert or ensure this is the smallest fragment index inserting this key + is_unique = map_insert_ref.insert_or_apply( + slot_type{static_cast(val_idx), frag_idx}, cuco::reduce::min{}); uniq_elem_size = [&]() -> size_type { if (not is_unique) { return 0; } switch (col->physical_type) { @@ -175,13 +172,11 @@ struct map_insert_fn { auto num_unique = block_reduce(reduce_storage).Sum(is_unique); __syncthreads(); auto uniq_data_size = block_reduce(reduce_storage).Sum(uniq_elem_size); - // One thread atomically updates the number and data size of total unique values as well as - // the number of unique values inserted by this fragment. + // One thread atomically updates the number and data size of total unique values. if (t == 0) { total_num_dict_entries = chunk_num_dict_entries.fetch_add(num_unique, cuda::std::memory_order_relaxed); total_num_dict_entries += num_unique; - num_dict_vals += num_unique; chunk_uniq_data_size.fetch_add(uniq_data_size, cuda::std::memory_order_relaxed); } __syncthreads(); @@ -189,8 +184,6 @@ struct map_insert_fn { // Check if the num unique values in chunk has already exceeded max dict size and early exit if (total_num_dict_entries > MAX_DICT_SIZE) { break; } } // for loop - // Flush the number of unique values inserted by this fragment - if (t == 0) { frag->num_dict_vals = num_dict_vals; }; } else { CUDF_UNREACHABLE("Unsupported type to insert in map"); } @@ -317,24 +310,40 @@ CUDF_KERNEL void __launch_bounds__(block_size) auto const& col_frags = frags[col_idx]; auto const frag_start = static_cast(chunk.fragments - col_frags.data()); - // Initialize fragment_offsets with prefix sum (exclusive) of number of dictionary values - // inserted by each fragment in this chunk. + // fragment_offsets will contain the prefix-sum of number of dictionary values first seen in + // each page fragment __shared__ size_type fragment_offsets[MAX_FRAGMENTS_PER_CHUNK]; + + // Initialize all fragment offsets to 0 + for (auto idx = t; idx < num_frags; idx += block_size) { + fragment_offsets[idx] = 0; + } + __syncthreads(); + + // Iterate over slots and count the number of dict values first seen page fragment { - using block_scan = cub::BlockScan; - __shared__ typename block_scan::TempStorage scan_storage; - - auto base_idx = 0; - while (base_idx < num_frags) { - auto const idx = base_idx + t; - auto const per_thread_count = - (idx < num_frags) ? col_frags[frag_start + idx].num_dict_vals : 0; - auto per_thread_offset = 0; - block_scan(scan_storage).ExclusiveSum(per_thread_count, per_thread_offset); - if (idx < num_frags) { fragment_offsets[idx] = per_thread_offset; } - base_idx += block_size; - __syncthreads(); + for (auto slot_idx = t; slot_idx < chunk.dict_map_size; slot_idx += block_size) { + auto const* slot = map_storage.data() + chunk.dict_map_offset + slot_idx; + if (slot->first != KEY_SENTINEL) { + auto const frag_loc = static_cast(slot->second) - frag_start; + cudf_assert(frag_loc >= 0 && frag_loc < num_frags && + "fragment index in the slot is out of range of the chunk"); + atomicAdd(&fragment_offsets[frag_loc], 1); + } } + __syncthreads(); + } + + // Exclusive scan to convert counts to offsets + using block_scan = cub::BlockScan; + __shared__ typename block_scan::TempStorage scan_storage; + + { + auto const per_thread_count = (t < num_frags) ? fragment_offsets[t] : 0; + auto per_thread_offset = 0; + block_scan(scan_storage).ExclusiveSum(per_thread_count, per_thread_offset); + if (t < num_frags) { fragment_offsets[t] = per_thread_offset; } + __syncthreads(); } // Iterate over slots and claim a dictionary index from the fragment offsets (spatial-locality @@ -344,9 +353,7 @@ CUDF_KERNEL void __launch_bounds__(block_size) auto const key = slot->first; if (key != KEY_SENTINEL) { auto const frag_loc = static_cast(slot->second) - frag_start; - cudf_assert(frag_loc >= 0 && frag_loc < num_frags && - "populate stamped a fragment hint outside this chunk's fragment range"); - auto const loc = atomicAdd(&fragment_offsets[frag_loc], 1); + auto const loc = atomicAdd(&fragment_offsets[frag_loc], 1); cudf_assert(loc < MAX_DICT_SIZE && "Number of filled slots exceeds max dict size"); chunk.dict_data[loc] = key; slot->second = loc; diff --git a/cpp/src/io/parquet/page_enc.cu b/cpp/src/io/parquet/page_enc.cu index 4dae7dac6c98..13792bcca95b 100644 --- a/cpp/src/io/parquet/page_enc.cu +++ b/cpp/src/io/parquet/page_enc.cu @@ -153,7 +153,6 @@ void __device__ init_frag_state(frag_init_state_s* const s, // smaller. num_rows is fixed but fragment size could be larger if the data is strings or // nested. s->frag.num_rows = min(fragment_size, part_end_row - s->frag.start_row); - s->frag.num_dict_vals = 0; s->frag.fragment_data_size = 0; s->frag.dict_data_size = 0; diff --git a/cpp/src/io/parquet/parquet_gpu.hpp b/cpp/src/io/parquet/parquet_gpu.hpp index b594618ca72a..3a177b7b62e6 100644 --- a/cpp/src/io/parquet/parquet_gpu.hpp +++ b/cpp/src/io/parquet/parquet_gpu.hpp @@ -530,7 +530,6 @@ struct PageFragment { uint32_t num_valid; // Date: Fri, 15 May 2026 23:48:01 +0000 Subject: [PATCH 26/28] Remove unnecessary changes from #22279 --- cpp/src/io/parquet/chunk_dict.cu | 25 ++++++++++++------------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index 425270e78415..61f1b5846b0c 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -99,9 +99,6 @@ struct map_insert_fn { auto const col = chunk->col_desc; column_device_view const& data_col = *col->leaf_column; __shared__ size_type total_num_dict_entries; - __shared__ size_type num_dict_vals; - if (t == 0) { num_dict_vals = 0; }; - __syncthreads(); using equality_fn_type = equality_functor; using hash_fn_type = hash_functor; @@ -175,21 +172,18 @@ struct map_insert_fn { auto num_unique = block_reduce(reduce_storage).Sum(is_unique); __syncthreads(); auto uniq_data_size = block_reduce(reduce_storage).Sum(uniq_elem_size); - // One thread atomically updates the number and data size of total unique values. + // First thread atomically updates the total number and data size of unique values if (t == 0) { total_num_dict_entries = chunk_num_dict_entries.fetch_add(num_unique, cuda::std::memory_order_relaxed); total_num_dict_entries += num_unique; - num_dict_vals += num_unique; chunk_uniq_data_size.fetch_add(uniq_data_size, cuda::std::memory_order_relaxed); } __syncthreads(); // Check if the num unique values in chunk has already exceeded max dict size and early exit - if (total_num_dict_entries > MAX_DICT_SIZE) { break; } + if (total_num_dict_entries > MAX_DICT_SIZE) { return; } } // for loop - // Flush the number of unique values inserted by this fragment - if (t == 0) { frag->num_dict_vals = num_dict_vals; }; } else { CUDF_UNREACHABLE("Unsupported type to insert in map"); } @@ -345,11 +339,16 @@ CUDF_KERNEL void __launch_bounds__(block_size) __shared__ typename block_scan::TempStorage scan_storage; { - auto const per_thread_count = (t < num_frags) ? fragment_offsets[t] : 0; - auto per_thread_offset = 0; - block_scan(scan_storage).ExclusiveSum(per_thread_count, per_thread_offset); - if (t < num_frags) { fragment_offsets[t] = per_thread_offset; } - __syncthreads(); + auto base_idx = uint32_t{0}; + while (base_idx < num_frags) { + auto const idx = base_idx + t; + auto const per_thread_count = (idx < num_frags) ? fragment_offsets[t] : 0; + auto per_thread_offset = 0; + block_scan(scan_storage).ExclusiveSum(per_thread_count, per_thread_offset); + if (idx < num_frags) { fragment_offsets[idx] = per_thread_offset; } + base_idx += block_size; + __syncthreads(); + } } // Iterate over slots and claim a dictionary index from the fragment offsets (spatial-locality From 30d3ae9073ca0c9f3af33f8b2dabd95f4e903ddd Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Fri, 15 May 2026 23:51:08 +0000 Subject: [PATCH 27/28] Minor improvements --- cpp/src/io/parquet/chunk_dict.cu | 4 +++- cpp/src/io/parquet/page_enc.cu | 2 +- cpp/tests/io/parquet_writer_test.cpp | 12 +++++------- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/cpp/src/io/parquet/chunk_dict.cu b/cpp/src/io/parquet/chunk_dict.cu index 61f1b5846b0c..61feffb6d5e2 100644 --- a/cpp/src/io/parquet/chunk_dict.cu +++ b/cpp/src/io/parquet/chunk_dict.cu @@ -358,7 +358,9 @@ CUDF_KERNEL void __launch_bounds__(block_size) auto const key = slot->first; if (key != KEY_SENTINEL) { auto const frag_loc = static_cast(slot->second) - frag_start; - auto const loc = atomicAdd(&fragment_offsets[frag_loc], 1); + cudf_assert(frag_loc >= 0 && frag_loc < num_frags && + "populate stamped a fragment hint outside this chunk's fragment range"); + auto const loc = atomicAdd(&fragment_offsets[frag_loc], 1); cudf_assert(loc < MAX_DICT_SIZE && "Number of filled slots exceeds max dict size"); chunk.dict_data[loc] = key; slot->second = loc; diff --git a/cpp/src/io/parquet/page_enc.cu b/cpp/src/io/parquet/page_enc.cu index 13792bcca95b..2e2f237a6331 100644 --- a/cpp/src/io/parquet/page_enc.cu +++ b/cpp/src/io/parquet/page_enc.cu @@ -152,7 +152,7 @@ void __device__ init_frag_state(frag_init_state_s* const s, // frag.num_rows = fragment_size except for the last fragment in partition which can be // smaller. num_rows is fixed but fragment size could be larger if the data is strings or // nested. - s->frag.num_rows = min(fragment_size, part_end_row - s->frag.start_row); + s->frag.num_rows = cuda::std::min(fragment_size, part_end_row - s->frag.start_row); s->frag.fragment_data_size = 0; s->frag.dict_data_size = 0; diff --git a/cpp/tests/io/parquet_writer_test.cpp b/cpp/tests/io/parquet_writer_test.cpp index 64bfd2974d24..859059ea4ff8 100644 --- a/cpp/tests/io/parquet_writer_test.cpp +++ b/cpp/tests/io/parquet_writer_test.cpp @@ -1197,13 +1197,11 @@ TEST_F(ParquetWriterTest, VariableBitWidthDictEncoding) EXPECT_LE(*max_bits_iter, chunk_wide_max_bits); // Check expected number of freq and rare pages - - auto const total_page_count = static_cast(page_dict_bits.size()); - auto const freq_page_count = static_cast( - std::ranges::count_if(page_dict_bits, [&](int nbits) { return nbits <= frequent_max_bits; - })); - EXPECT_EQ(freq_page_count, freq_pages); - EXPECT_EQ(total_page_count - freq_page_count, num_pages - freq_pages); + auto const total_page_count = static_cast(page_dict_bits.size()); + auto const freq_page_count = static_cast( + std::ranges::count_if(page_dict_bits, [&](int nbits) { return nbits <= frequent_max_bits; })); + EXPECT_EQ(freq_page_count, freq_pages); + EXPECT_EQ(total_page_count - freq_page_count, num_pages - freq_pages); } } From 21898fde3bdd3c978788a2de668c675a1ab99e70 Mon Sep 17 00:00:00 2001 From: Muhammad Haseeb Date: Mon, 1 Jun 2026 22:43:48 +0000 Subject: [PATCH 28/28] style fix --- cpp/src/io/parquet/parquet_gpu.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/io/parquet/parquet_gpu.hpp b/cpp/src/io/parquet/parquet_gpu.hpp index 4b9595c90403..7a451a2b475b 100644 --- a/cpp/src/io/parquet/parquet_gpu.hpp +++ b/cpp/src/io/parquet/parquet_gpu.hpp @@ -530,7 +530,7 @@ struct PageFragment { //!< non-leaf level uint32_t num_valid; //