From 6195c4367834f30a73d04169c88a114f782cd9a8 Mon Sep 17 00:00:00 2001 From: Akihiro Hirota Date: Sun, 6 Sep 2026 19:50:20 +0900 Subject: [PATCH 1/2] Fix Parquet V2 inputs to decompression scratch queries --- .../io/parquet/parquet_reader_chunks.cpp | 13 +- .../io/parquet/reader_impl_chunking_utils.cu | 15 +- cpp/tests/io/parquet_chunked_reader_test.cu | 169 ++++++++++++++++++ 3 files changed, 189 insertions(+), 8 deletions(-) diff --git a/cpp/benchmarks/io/parquet/parquet_reader_chunks.cpp b/cpp/benchmarks/io/parquet/parquet_reader_chunks.cpp index 9d58cce16b70..b1b0b440c834 100644 --- a/cpp/benchmarks/io/parquet/parquet_reader_chunks.cpp +++ b/cpp/benchmarks/io/parquet/parquet_reader_chunks.cpp @@ -86,9 +86,10 @@ void BM_parquet_read_subrowgroup_chunks(nvbench::state& state, auto const pass_read_limit = static_cast(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(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 = [&]() { @@ -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(v2_headers); 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); @@ -160,6 +163,8 @@ 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", "ZSTD"}) + .add_int64_axis("write_v2_headers", {0, 1}) .set_min_samples(4) .add_int64_axis("cardinality", {0, 1000}) .add_int64_axis("run_length", {1, 32}) diff --git a/cpp/src/io/parquet/reader_impl_chunking_utils.cu b/cpp/src/io/parquet/reader_impl_chunking_utils.cu index 68c9ff5f6581..4ee08333f53a 100644 --- a/cpp/src/io/parquet/reader_impl_chunking_utils.cu +++ b/cpp/src/io/parquet/reader_impl_chunking_utils.cu @@ -768,11 +768,18 @@ rmm::device_uvector compute_decompression_scratch_sizes( 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( - page.page_data, static_cast(page.compressed_page_size)); + // Match the inputs passed to decompression: V2 levels are not compressed, and V2 + // pages may contain uncompressed values or no values after the levels. + auto const is_page_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 (parquet_compression_support(chunks[page.chunk_idx].codec).first == codec and + is_page_compressed and page.compressed_page_size > offset) { + temp_spans[i] = {page.page_data + offset, + static_cast(page.compressed_page_size - offset)}; } else { - temp_spans[i] = device_span(); // Mark pages with other codecs as empty + temp_spans[i] = {}; } }); // Copy only non-null spans diff --git a/cpp/tests/io/parquet_chunked_reader_test.cu b/cpp/tests/io/parquet_chunked_reader_test.cu index 7abb3cf527af..02c74689c98f 100644 --- a/cpp/tests/io/parquet_chunked_reader_test.cu +++ b/cpp/tests/io/parquet_chunked_reader_test.cu @@ -26,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -42,6 +43,7 @@ #include #include #include +#include #include #include @@ -200,6 +202,173 @@ struct ParquetChunkedReaderTest : public cudf::test::BaseFixture {}; using ParquetChunkedDecompressionTest = DecompressionTest; +// Exercise has_next() before the first read, including the pre-parsed-footer entry point. +struct ParquetScratchTest : public ParquetChunkedReaderTest { + void check_chunked_read(std::string const& filepath, cudf::table_view expected, bool use_metadata) + { + auto const options = + cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}).build(); + auto const full = cudf::io::read_parquet(options); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected, full.tbl->view()); + for (auto const limit : {std::size_t{32'768}, std::size_t{1'024'000'000}}) { + auto reader = [&] { + if (use_metadata) { + auto sources = cudf::io::make_datasources(cudf::io::source_info{filepath}); + auto footers = cudf::io::read_parquet_footers(sources); + return cudf::io::chunked_parquet_reader( + 0, limit, std::move(sources), std::move(footers), options); + } + return cudf::io::chunked_parquet_reader(0, limit, options); + }(); + ASSERT_TRUE(reader.has_next()); + auto const [result, num_chunks] = chunked_read(reader); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected, result->view()); + EXPECT_GT(num_chunks, 0); + } + } +}; + +struct ParquetV2ScratchLevelsTest + : public ParquetScratchTest, + public ::testing::WithParamInterface> {}; + +TEST_P(ParquetV2ScratchLevelsTest, NullableAndRepeatedValues) +{ + auto const [v2, dictionary, use_metadata] = GetParam(); + tmp_env_var const nvcomp{nvcomp_policy_env_var, "ALWAYS"}; + tmp_env_var const host_decomp{host_decomp_env_var, "OFF"}; + constexpr cudf::size_type num_rows = 20'000; + auto values = cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i % 25; }); + auto valid = + cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i % 10 != 0; }); + int64s_col flat(values, values + num_rows, valid); + int32s_col child(values, values + 2 * num_rows, valid); + auto offsets = cudf::detail::make_counting_transform_iterator(0, [](auto i) { return 2 * i; }); + int32s_col list_offsets(offsets, offsets + num_rows + 1); + auto const lists = cudf::make_lists_column( + num_rows, list_offsets.release(), child.release(), 0, rmm::device_buffer{}); + auto const expected = cudf::table_view{{flat, lists->view()}}; + auto metadata = cudf::io::table_input_metadata{expected}; + if (not dictionary) { + metadata.column_metadata[0].set_encoding(cudf::io::column_encoding::PLAIN); + metadata.column_metadata[1].child(1).set_encoding(cudf::io::column_encoding::PLAIN); + } + auto const filepath = temp_env->get_temp_filepath("ScratchLevels.parquet"); + auto const options = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, expected) + .metadata(metadata) + .compression(cudf::io::compression_type::ZSTD) + .write_v2_headers(v2) + .page_level_compression(v2) + .dictionary_policy(dictionary ? cudf::io::dictionary_policy::ALWAYS + : cudf::io::dictionary_policy::NEVER) + .max_page_size_rows(5'000) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .build(); + cudf::io::write_parquet(options); + + auto const source = cudf::io::datasource::create(filepath); + cudf::io::parquet::FileMetaData footer; + read_footer(source, &footer); + ASSERT_EQ(footer.row_groups.size(), 1); + ASSERT_EQ(footer.row_groups.front().columns.size(), 2); + for (auto i = 0; i < 2; ++i) { + auto const& column = footer.row_groups.front().columns[i]; + ASSERT_EQ(column.meta_data.codec, cudf::io::parquet::Compression::ZSTD); + EXPECT_EQ(column.meta_data.dictionary_page_offset > 0, dictionary); + auto const indexes = read_offset_index(source, column); + ASSERT_GT(indexes.page_locations.size(), 1); + for (auto const& location : indexes.page_locations) { + auto const page = read_page_header(source, location); + ASSERT_EQ( + page.type, + v2 ? cudf::io::parquet::PageType::DATA_PAGE_V2 : cudf::io::parquet::PageType::DATA_PAGE); + if (v2) { + EXPECT_GT(page.data_page_header_v2.definition_levels_byte_length, 0); + if (i == 1) { EXPECT_GT(page.data_page_header_v2.repetition_levels_byte_length, 0); } + EXPECT_TRUE(page.data_page_header_v2.is_compressed); + } + } + } + check_chunked_read(filepath, expected, use_metadata); +} + +INSTANTIATE_TEST_SUITE_P(PageVersionAndReader, + ParquetV2ScratchLevelsTest, + ::testing::Combine(::testing::Bool(), + ::testing::Bool(), + ::testing::Bool())); + +struct ParquetV2ScratchSkippedPagesTest + : public ParquetScratchTest, + public ::testing::WithParamInterface> {}; + +TEST_P(ParquetV2ScratchSkippedPagesTest, UncompressedAndEmptyValues) +{ + auto const [empty_pages, use_metadata] = GetParam(); + tmp_env_var const nvcomp{nvcomp_policy_env_var, "ALWAYS"}; + tmp_env_var const host_decomp{host_decomp_env_var, "OFF"}; + constexpr cudf::size_type num_rows = 20'000; + std::vector values(num_rows, 42); + std::mt19937_64 random(42); + std::generate(values.begin(), values.begin() + num_rows / 2, [&] { return random(); }); + // Keep compressed values in the same column chunk so empty pages are still tagged ZSTD. + auto valid = cudf::detail::make_counting_transform_iterator( + 0, [empty_pages](auto i) { return not empty_pages or i >= num_rows / 2; }); + int64s_col column(values.begin(), values.end(), valid); + auto const expected = cudf::table_view{{column}}; + auto metadata = cudf::io::table_input_metadata{expected}; + metadata.column_metadata.front().set_encoding(cudf::io::column_encoding::PLAIN); + auto const filepath = temp_env->get_temp_filepath("ScratchSkippedPages.parquet"); + auto const options = + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, expected) + .metadata(metadata) + .compression(cudf::io::compression_type::ZSTD) + .write_v2_headers(true) + .page_level_compression(true) + .dictionary_policy(cudf::io::dictionary_policy::NEVER) + .max_page_size_rows(5'000) + .stats_level(cudf::io::statistics_freq::STATISTICS_COLUMN) + .build(); + cudf::io::write_parquet(options); + + auto const source = cudf::io::datasource::create(filepath); + cudf::io::parquet::FileMetaData footer; + read_footer(source, &footer); + ASSERT_EQ(footer.row_groups.size(), 1); + ASSERT_EQ(footer.row_groups.front().columns.size(), 1); + auto const& chunk = footer.row_groups.front().columns.front(); + ASSERT_EQ(chunk.meta_data.codec, cudf::io::parquet::Compression::ZSTD); + auto const indexes = read_offset_index(source, chunk); + ASSERT_GT(indexes.page_locations.size(), 1); + bool found_compressed = false; + bool found_uncompressed = false; + bool found_empty = false; + for (auto const& location : indexes.page_locations) { + auto const page = read_page_header(source, location); + ASSERT_EQ(page.type, cudf::io::parquet::PageType::DATA_PAGE_V2); + auto const& header = page.data_page_header_v2; + found_compressed |= header.is_compressed; + found_uncompressed |= not header.is_compressed; + if (empty_pages and header.num_nulls == header.num_values) { + found_empty = true; + EXPECT_EQ(page.compressed_page_size, + header.definition_levels_byte_length + header.repetition_levels_byte_length); + } + } + EXPECT_TRUE(found_compressed); + if (empty_pages) { + EXPECT_TRUE(found_empty); + } else { + EXPECT_TRUE(found_uncompressed); + } + check_chunked_read(filepath, expected, use_metadata); +} + +INSTANTIATE_TEST_SUITE_P(PageContentsAndReader, + ParquetV2ScratchSkippedPagesTest, + ::testing::Combine(::testing::Bool(), ::testing::Bool())); + TEST_F(ParquetChunkedReaderTest, TestChunkedReadNoData) { std::vector> input_columns; From 2fe0ccba4552f887748b8247c80ff86180092f94 Mon Sep 17 00:00:00 2001 From: Akihiro Hirota Date: Mon, 7 Sep 2026 03:28:23 +0900 Subject: [PATCH 2/2] Align Parquet scratch accounting and strengthen regression coverage --- .../io/parquet/parquet_reader_chunks.cpp | 26 +++- .../io/parquet/reader_impl_chunking_utils.cu | 40 +++--- .../io/parquet/reader_impl_chunking_utils.cuh | 21 ++- cpp/tests/io/parquet_chunked_reader_test.cu | 122 +++++++++++++++++- 4 files changed, 175 insertions(+), 34 deletions(-) diff --git a/cpp/benchmarks/io/parquet/parquet_reader_chunks.cpp b/cpp/benchmarks/io/parquet/parquet_reader_chunks.cpp index b1b0b440c834..8f264482bedb 100644 --- a/cpp/benchmarks/io/parquet/parquet_reader_chunks.cpp +++ b/cpp/benchmarks/io/parquet/parquet_reader_chunks.cpp @@ -103,7 +103,7 @@ void BM_parquet_read_subrowgroup_chunks(nvbench::state& state, cudf::io::parquet_writer_options::builder(source_sink.make_sink_info(), view) .compression(compression) .write_v2_headers(v2_headers) - .page_level_compression(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); @@ -163,8 +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", "ZSTD"}) - .add_int64_axis("write_v2_headers", {0, 1}) + .add_string_axis("compression_type", {"SNAPPY"}) + .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}) @@ -173,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; + +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}); diff --git a/cpp/src/io/parquet/reader_impl_chunking_utils.cu b/cpp/src/io/parquet/reader_impl_chunking_utils.cu index 4ee08333f53a..499516d922da 100644 --- a/cpp/src/io/parquet/reader_impl_chunking_utils.cu +++ b/cpp/src/io/parquet/reader_impl_chunking_utils.cu @@ -732,7 +732,7 @@ rmm::device_uvector compute_decompression_scratch_sizes( auto temp_cost = cudf::detail::make_pinned_vector_async(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 d_temp_cost = @@ -756,32 +756,22 @@ rmm::device_uvector 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> 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]; - // Match the inputs passed to decompression: V2 levels are not compressed, and V2 - // pages may contain uncompressed values or no values after the levels. - auto const is_page_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 (parquet_compression_support(chunks[page.chunk_idx].codec).first == codec and - is_page_compressed and page.compressed_page_size > offset) { - temp_spans[i] = {page.page_data + offset, - static_cast(page.compressed_page_size - offset)}; - } else { - temp_spans[i] = {}; - } - }); + 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{}; + }); // Copy only non-null spans rmm::device_uvector> page_spans(pages.size(), stream); auto end_iter = diff --git a/cpp/src/io/parquet/reader_impl_chunking_utils.cuh b/cpp/src/io/parquet/reader_impl_chunking_utils.cuh index 6536d3d7d4f9..e6a8ce296494 100644 --- a/cpp/src/io/parquet/reader_impl_chunking_utils.cuh +++ b/cpp/src/io/parquet/reader_impl_chunking_utils.cuh @@ -217,7 +217,7 @@ void detect_malformed_pages(device_span pages, /** * @brief Computes the per-page scratch space required for decompression. */ -rmm::device_uvector compute_decompression_scratch_sizes( +CUDF_EXPORT rmm::device_uvector compute_decompression_scratch_sizes( device_span chunks, device_span pages, cuda::stream_ref stream); @@ -416,6 +416,20 @@ struct codec_stats { host_span page_mask); }; +/** + * @brief Returns the compressed values passed to the decompressor, excluding V2 level bytes. + */ +struct get_decompression_input { + __device__ inline device_span 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(page.compressed_page_size - offset)}; + } +}; + /** * @brief Functor which retrieves per-page decompression information. */ @@ -423,7 +437,10 @@ struct get_decomp_info { device_span 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}; } + // Keep the full page size as a conservative bound, as in codec_stats::add_pages. + return {codec, 1, static_cast(p.uncompressed_page_size), static_cast(p.uncompressed_page_size)}; diff --git a/cpp/tests/io/parquet_chunked_reader_test.cu b/cpp/tests/io/parquet_chunked_reader_test.cu index 02c74689c98f..a697cfbd7dff 100644 --- a/cpp/tests/io/parquet_chunked_reader_test.cu +++ b/cpp/tests/io/parquet_chunked_reader_test.cu @@ -4,6 +4,7 @@ */ #include "compression_common.hpp" +#include "io/parquet/reader_impl_chunking_utils.cuh" #include "io_test_utils.hpp" #include "parquet_common.hpp" @@ -22,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -38,6 +40,7 @@ #include #include +#include #include #include @@ -228,13 +231,122 @@ struct ParquetScratchTest : public ParquetChunkedReaderTest { } }; +TEST_F(ParquetScratchTest, DecompressionInputs) +{ + namespace pq = cudf::io::parquet::detail; + auto const stream = cudf::get_default_stream(); + rmm::device_uvector data(100, stream); + + // V1 and dictionary pages ignore is_compressed; V2 may have either/both level streams, + // uncompressed values, or no values. No compressed input is dereferenced in this test. + std::vector pages(7); + for (auto& page : pages) { + page.page_data = data.data(); + page.compressed_page_size = 100; + page.uncompressed_page_size = 1'000; + } + pages[1].flags = pq::PAGEINFO_FLAGS_DICTIONARY; + for (auto i = 2; i < 7; ++i) { + pages[i].flags = pq::PAGEINFO_FLAGS_V2; + pages[i].is_compressed = true; + } + pages[2].lvl_bytes[pq::level_type::DEFINITION] = 10; + pages[3].lvl_bytes[pq::level_type::REPETITION] = 20; + pages[4].lvl_bytes[pq::level_type::DEFINITION] = 10; + pages[4].lvl_bytes[pq::level_type::REPETITION] = 20; + pages[5].is_compressed = false; + pages[6].lvl_bytes[pq::level_type::DEFINITION] = 100; + auto const d_pages = + cudf::detail::make_device_uvector(pages, stream, cudf::get_current_device_resource_ref()); + rmm::device_uvector> inputs(pages.size(), stream); + thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + d_pages.begin(), + d_pages.end(), + inputs.begin(), + pq::get_decompression_input{}); + auto const h_inputs = cudf::detail::make_host_vector(inputs, stream); + std::vector const offsets{0, 0, 10, 20, 30}; + for (std::size_t i = 0; i < pages.size(); ++i) { + SCOPED_TRACE(i); + bool const compressed = i < offsets.size(); + EXPECT_EQ(h_inputs[i].data(), compressed ? data.data() + offsets[i] : nullptr); + EXPECT_EQ(h_inputs[i].size(), compressed ? 100 - offsets[i] : 0); + } +} + +TEST_F(ParquetScratchTest, SkippedPagesDoNotReduceScratchEstimate) +{ + namespace pq = cudf::io::parquet::detail; + tmp_env_var const nvcomp{nvcomp_policy_env_var, "ALWAYS"}; + tmp_env_var const host_decomp{host_decomp_env_var, "OFF"}; + auto const stream = cudf::get_default_stream(); + // A Zstd frame containing 32,768 repetitions of '*'. + std::vector const frame{0x28, + 0xb5, + 0x2f, + 0xfd, + 0x60, + 0x00, + 0x7f, + 0x4d, + 0x00, + 0x00, + 0x10, + 0x2a, + 0x2a, + 0x01, + 0x00, + 0xfb, + 0xff, + 0x0e, + 0xb0}; + auto data = + cudf::detail::make_device_uvector(frame, stream, cudf::get_current_device_resource_ref()); + std::vector chunks(1); + chunks.front().codec = cudf::io::parquet::Compression::ZSTD; + auto const d_chunks = + cudf::detail::make_device_uvector(chunks, stream, cudf::get_current_device_resource_ref()); + pq::PageInfo compressed{}; + compressed.page_data = data.data(); + compressed.flags = pq::PAGEINFO_FLAGS_V2; + compressed.is_compressed = true; + compressed.compressed_page_size = frame.size(); + compressed.uncompressed_page_size = 32'768; + auto uncompressed = compressed; + uncompressed.is_compressed = false; + auto empty = compressed; + empty.compressed_page_size = 0; + empty.uncompressed_page_size = 0; + auto const scratch_sizes = [&](std::vector const& pages) { + auto const d_pages = + cudf::detail::make_device_uvector(pages, stream, cudf::get_current_device_resource_ref()); + return cudf::detail::make_host_vector( + pq::compute_decompression_scratch_sizes(d_chunks, d_pages, stream), stream); + }; + auto const baseline = scratch_sizes({compressed, compressed}); + auto const mixed = + scratch_sizes({compressed, uncompressed, empty, compressed, uncompressed, empty}); + auto const skipped = scratch_sizes({uncompressed, empty}); + ASSERT_GT(baseline.front(), 0); + // Skipped pages must neither add scratch nor dilute the extended/legacy adjustment ratio. + EXPECT_EQ(mixed[0], baseline[0]); + EXPECT_EQ(mixed[1], baseline[0]); + EXPECT_EQ(mixed[2], baseline[0]); + EXPECT_EQ(mixed[3], baseline[1]); + EXPECT_EQ(mixed[4], baseline[1]); + EXPECT_EQ(mixed[5], baseline[1]); + EXPECT_EQ(skipped[0], 0); + EXPECT_EQ(skipped[1], 0); +} + struct ParquetV2ScratchLevelsTest : public ParquetScratchTest, - public ::testing::WithParamInterface> {}; + public ::testing::WithParamInterface, bool, bool>> {}; TEST_P(ParquetV2ScratchLevelsTest, NullableAndRepeatedValues) { - auto const [v2, dictionary, use_metadata] = GetParam(); + auto const [page_options, dictionary, use_metadata] = GetParam(); + auto const [v2, page_compression] = page_options; tmp_env_var const nvcomp{nvcomp_policy_env_var, "ALWAYS"}; tmp_env_var const host_decomp{host_decomp_env_var, "OFF"}; constexpr cudf::size_type num_rows = 20'000; @@ -259,7 +371,7 @@ TEST_P(ParquetV2ScratchLevelsTest, NullableAndRepeatedValues) .metadata(metadata) .compression(cudf::io::compression_type::ZSTD) .write_v2_headers(v2) - .page_level_compression(v2) + .page_level_compression(page_compression) .dictionary_policy(dictionary ? cudf::io::dictionary_policy::ALWAYS : cudf::io::dictionary_policy::NEVER) .max_page_size_rows(5'000) @@ -295,7 +407,9 @@ TEST_P(ParquetV2ScratchLevelsTest, NullableAndRepeatedValues) INSTANTIATE_TEST_SUITE_P(PageVersionAndReader, ParquetV2ScratchLevelsTest, - ::testing::Combine(::testing::Bool(), + ::testing::Combine(::testing::Values(std::pair{false, false}, + std::pair{true, false}, + std::pair{true, true}), ::testing::Bool(), ::testing::Bool()));