Skip to content

Fix Parquet V2 inputs to decompression scratch queries #24006

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
a-hirota wants to merge 2 commits into
base: main
Choose a base branch
from
+346 −24
Open

Fix Parquet V2 inputs to decompression scratch queries #24006

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6195c43
Fix Parquet V2 inputs to decompression scratch queries
a-hirota Sep 6, 2026
2fe0ccb
Align Parquet scratch accounting and strengthen regression coverage
a-hirota Sep 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 29 additions & 4 deletions cpp/benchmarks/io/parquet/parquet_reader_chunks.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,9 +86,10 @@ void BM_parquet_read_subrowgroup_chunks(nvbench::state& state,
auto const pass_read_limit = static_cast<cudf::size_type>(state.get_int64("pass_read_limit"));
auto const source_type = retrieve_io_type_enum(state.get_string("io_type"));
auto const data_size = static_cast<size_t>(state.get_int64("data_size"));
auto const compression = cudf::io::compression_type::SNAPPY;
auto const rg_size_bytes = state.get_int64("row_group_size_bytes");
auto const rg_size_rows = state.get_int64("row_group_size_rows");
auto const compression = retrieve_compression_type_enum(state.get_string("compression_type"));
auto const v2_headers = state.get_int64("write_v2_headers") != 0;
auto const rg_size_bytes = state.get_int64("row_group_size_bytes");
auto const rg_size_rows = state.get_int64("row_group_size_rows");
cuio_source_sink_pair source_sink(source_type);

auto const num_rows_written = [&]() {
Expand All @@ -100,7 +101,9 @@ void BM_parquet_read_subrowgroup_chunks(nvbench::state& state,

cudf::io::parquet_writer_options write_opts =
cudf::io::parquet_writer_options::builder(source_sink.make_sink_info(), view)
.compression(compression);
.compression(compression)
.write_v2_headers(v2_headers)
.page_level_compression(state.get_int64("page_level_compression") != 0);
if (rg_size_bytes > 0) write_opts.set_row_group_size_bytes(rg_size_bytes);
if (rg_size_rows > 0) write_opts.set_row_group_size_rows(rg_size_rows);
cudf::io::write_parquet(write_opts);
Expand Down Expand Up @@ -160,6 +163,9 @@ NVBENCH_BENCH_TYPES(BM_parquet_read_chunks, NVBENCH_TYPE_AXES(d_type_list))
NVBENCH_BENCH_TYPES(BM_parquet_read_subrowgroup_chunks, NVBENCH_TYPE_AXES(d_type_list))
.set_name("parquet_read_subrowgroup_chunks")
.add_string_axis("io_type", {"DEVICE_BUFFER"})
.add_string_axis("compression_type", {"SNAPPY"})

@vuule vuule Sep 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Adding compression_type, write_v2_headers, and page_level_compression as single-value axes is forced by sharing one callable across two registrations, but it changes every existing parquet_read_subrowgroup_chunks state string and so breaks continuity with previously recorded results. A thin wrapper function that hardcodes the Snappy/V1 values for the original benchmark would keep the old axis set intact.

.add_int64_axis("write_v2_headers", {0})
.add_int64_axis("page_level_compression", {0})
.set_min_samples(4)
.add_int64_axis("cardinality", {0, 1000})
.add_int64_axis("run_length", {1, 32})
Expand All @@ -168,3 +174,22 @@ NVBENCH_BENCH_TYPES(BM_parquet_read_subrowgroup_chunks, NVBENCH_TYPE_AXES(d_type
.add_int64_axis("data_size", {512 << 20})
.add_int64_axis("row_group_size_bytes", {0})
.add_int64_axis("row_group_size_rows", {0});

// Keep the existing 512 MiB matrix unchanged; cover V2 scratch queries with 12 small cases.
using scratch_type_list =
nvbench::enum_type_list<data_type::INTEGRAL, data_type::STRING, data_type::LIST>;

NVBENCH_BENCH_TYPES(BM_parquet_read_subrowgroup_chunks, NVBENCH_TYPE_AXES(scratch_type_list))
.set_name("parquet_read_v2_scratch")
.add_string_axis("io_type", {"DEVICE_BUFFER"})
.add_string_axis("compression_type", {"ZSTD"})
.add_int64_axis("write_v2_headers", {1})
.add_int64_axis("page_level_compression", {0, 1})
.set_min_samples(4)
.add_int64_axis("cardinality", {1000})
.add_int64_axis("run_length", {1})
.add_int64_axis("chunk_read_limit", {0})
.add_int64_axis("pass_read_limit", {0, 500'000})
.add_int64_axis("data_size", {8 << 20})
.add_int64_axis("row_group_size_bytes", {0})
.add_int64_axis("row_group_size_rows", {0});
33 changes: 15 additions & 18 deletions cpp/src/io/parquet/reader_impl_chunking_utils.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -732,7 +732,7 @@ rmm::device_uvector<size_t> compute_decompression_scratch_sizes(
auto temp_cost = cudf::detail::make_pinned_vector_async<size_t>(pages.size(), stream);
auto h_decomp_info = cudf::detail::make_pinned_vector(decomp_info, stream);
std::transform(h_decomp_info.begin(), h_decomp_info.end(), temp_cost.begin(), [](auto const& d) {
return cudf::io::detail::get_decompression_scratch_size(d);
return d.num_pages == 0 ? 0 : cudf::io::detail::get_decompression_scratch_size(d);
});

rmm::device_uvector<size_t> d_temp_cost =
Expand All @@ -756,25 +756,22 @@ rmm::device_uvector<size_t> compute_decompression_scratch_sizes(
decomp_sum{},
stream);

// Collect pages with matching codecs
// Use the same page selection for the input spans and decompression-size estimates.
rmm::device_uvector<device_span<uint8_t const>> temp_spans(pages.size(), stream);
auto iter = cuda::counting_iterator{size_t{0}};
thrust::for_each(
rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
iter,
iter + pages.size(),
[pages = pages.begin(),
chunks = chunks.begin(),
temp_spans = temp_spans.begin(),
codec] __device__(size_t i) {
auto const& page = pages[i];
if (parquet_compression_support(chunks[page.chunk_idx].codec).first == codec) {
temp_spans[i] = device_span<uint8_t const>(
page.page_data, static_cast<size_t>(page.compressed_page_size));
} else {
temp_spans[i] = device_span<uint8_t const>(); // Mark pages with other codecs as empty
}
});
thrust::for_each(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()),
iter,
iter + pages.size(),
[pages = pages.begin(),
chunks = chunks.begin(),
temp_spans = temp_spans.begin(),
codec] __device__(size_t i) {
auto const& page = pages[i];
temp_spans[i] =
parquet_compression_support(chunks[page.chunk_idx].codec).first == codec
? get_decompression_input{}(page)
: device_span<uint8_t const>{};
});
// Copy only non-null spans

@vuule vuule Sep 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The copy_if predicate just below still keys off span.data() != nullptr. Now that empty spans are produced deliberately and carry meaning, not span.empty() expresses the intent directly and is robust if get_decompression_input ever returns a non-null zero-length span.

rmm::device_uvector<device_span<uint8_t const>> page_spans(pages.size(), stream);
auto end_iter =
Expand Down
21 changes: 19 additions & 2 deletions cpp/src/io/parquet/reader_impl_chunking_utils.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,7 +217,7 @@ void detect_malformed_pages(device_span<PageInfo const> pages,
/**
* @brief Computes the per-page scratch space required for decompression.
*/
rmm::device_uvector<size_t> compute_decompression_scratch_sizes(
CUDF_EXPORT rmm::device_uvector<size_t> compute_decompression_scratch_sizes(
device_span<ColumnChunkDesc const> chunks,
device_span<PageInfo const> pages,
cuda::stream_ref stream);
Expand Down Expand Up @@ -416,14 +416,31 @@ struct codec_stats {
host_span<bool const> page_mask);
};

/**
* @brief Returns the compressed values passed to the decompressor, excluding V2 level bytes.
*/
struct get_decompression_input {

@vuule vuule Sep 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The definition of "what the decompressor actually consumes" now lives in three places: codec_stats::add_pages (reader_impl_chunking_utils.cu:121-146), set_parameters (reader_impl_chunking_utils.cu:529-575), and this new functor. Divergence between the first two and the scratch query is precisely the bug being fixed here, so leaving three copies invites a repeat.

Declaring the functor CUDF_HOST_DEVICE and calling it from both host-side loops would collapse the duplication and make the invariant enforced by construction rather than by comment.

__device__ inline device_span<uint8_t const> operator()(PageInfo const& page) const
{
auto const is_compressed = (page.flags & PAGEINFO_FLAGS_V2) ? page.is_compressed : true;
auto const offset =
page.lvl_bytes[level_type::DEFINITION] + page.lvl_bytes[level_type::REPETITION];
if (not is_compressed or page.compressed_page_size <= offset) { return {}; }
return {page.page_data + offset, static_cast<size_t>(page.compressed_page_size - offset)};
}
};

/**
* @brief Functor which retrieves per-page decompression information.
*/
struct get_decomp_info {
device_span<ColumnChunkDesc const> chunks;
__device__ inline decompression_info operator()(PageInfo const& p) const
{
return {parquet_compression_support(chunks[p.chunk_idx].codec).first,
auto const codec = parquet_compression_support(chunks[p.chunk_idx].codec).first;
if (get_decompression_input{}(p).empty()) { return {codec, 0, 0, 0}; }

@vuule vuule Sep 8, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With get_decomp_info now excluding empty-payload V2 pages, the estimate can sit slightly below what decompression actually requests: codec_stats::add_pages still folds those pages' uncompressed_page_size into codec.total_decomp_size and max_decompressed_size, and both are passed straight to cudf::io::detail::decompress at reader_impl_chunking_utils.cu:618-619.

The gap is small (such pages carry only level bytes) and the previous code overestimated, so this is not a blocker, but it is the same class of divergence the PR sets out to remove and is worth either closing or calling out in the comment.

// Keep the full page size as a conservative bound, as in codec_stats::add_pages.
return {codec,
1,
static_cast<size_t>(p.uncompressed_page_size),
static_cast<size_t>(p.uncompressed_page_size)};
Expand Down
Loading
Loading