diff --git a/cpp/benchmarks/io/parquet/parquet_reader_chunks.cpp b/cpp/benchmarks/io/parquet/parquet_reader_chunks.cpp index 9d58cce16b70..8f264482bedb 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(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); @@ -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"}) + .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}) @@ -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; + +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 68c9ff5f6581..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,25 +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]; - 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)); - } else { - temp_spans[i] = device_span(); // 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{}; + }); // 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 7abb3cf527af..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,10 +23,12 @@ #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -37,11 +40,13 @@ #include #include +#include #include #include #include #include +#include #include #include @@ -200,6 +205,284 @@ 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); + } + } +}; + +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, bool, bool>> {}; + +TEST_P(ParquetV2ScratchLevelsTest, NullableAndRepeatedValues) +{ + 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; + 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(page_compression) + .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::Values(std::pair{false, false}, + std::pair{true, false}, + std::pair{true, true}), + ::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;