diff --git a/cpp/tests/io/parquet_chunked_reader_test.cu b/cpp/tests/io/parquet_chunked_reader_test.cu index 7abb3cf527a..01853f66dca 100644 --- a/cpp/tests/io/parquet_chunked_reader_test.cu +++ b/cpp/tests/io/parquet_chunked_reader_test.cu @@ -1270,20 +1270,21 @@ void input_limit_test_write(std::vector const& test_filenames, test_filenames[3], t, cudf::io::compression_type::SNAPPY, cudf::io::dictionary_policy::ALWAYS); } -void input_limit_test_read(std::vector const& test_filenames, - cudf::table_view const& t, - std::size_t output_limit, - std::size_t input_limit, - int const expected_chunk_counts[input_limit_expected_file_count]) +void input_limit_test_read( + std::vector const& test_filenames, + cudf::table_view const& t, + std::size_t output_limit, + std::size_t input_limit, + [[maybe_unused]] int const expected_chunk_counts[input_limit_expected_file_count], + bool require_multiple_chunks = false) { CUDF_EXPECTS(test_filenames.size() == input_limit_expected_file_count, "Unexpected count of test filenames"); for (std::size_t idx = 0; idx < test_filenames.size(); idx++) { auto result = chunked_read(test_filenames[idx], output_limit, input_limit); - // CUDF_EXPECTS(result.second == expected_chunk_counts[idx], - // "Unexpected number of chunks produced in chunk read"); CUDF_TEST_EXPECT_TABLES_EQUIVALENT(*result.first, t); + if (require_multiple_chunks) { EXPECT_GT(result.second, 1); } } } } // namespace @@ -1528,7 +1529,7 @@ TEST_F(ParquetChunkedReaderInputLimitTest, List) auto base_path = temp_env->get_temp_filepath("list"); auto test_filenames = input_limit_get_test_names(base_path); - constexpr int num_rows = 10'000'000; + constexpr int num_rows = 2'500'000; constexpr int list_size = 4; auto const stream = cudf::get_default_stream(); @@ -1575,16 +1576,13 @@ TEST_F(ParquetChunkedReaderInputLimitTest, List) // size of the decompressed data. so 2 GB is actually not enough to hold the whole thing at // once. // - // Note that in the dictionary cases, both of these revert down to 1 chunk because the - // dictionaries dramatically shrink the size of the uncompressed data. constexpr int expected_a[] = {3, 3, 1, 1}; - input_limit_test_read(test_filenames, tbl, 0, 256 * 1024 * 1024, expected_a); - // smaller limit + input_limit_test_read(test_filenames, tbl, 0, 64 * 1024 * 1024, expected_a); constexpr int expected_b[] = {5, 5, 2, 1}; - input_limit_test_read(test_filenames, tbl, 0, 128 * 1024 * 1024, expected_b); - // include output chunking as well + input_limit_test_read(test_filenames, tbl, 0, 32 * 1024 * 1024, expected_b); + // Include output chunking as well, and verify each input format is split. constexpr int expected_c[] = {10, 9, 8, 7}; - input_limit_test_read(test_filenames, tbl, 32 * 1024 * 1024, 64 * 1024 * 1024, expected_c); + input_limit_test_read(test_filenames, tbl, 8 * 1024 * 1024, 16 * 1024 * 1024, expected_c, true); } namespace { @@ -1678,7 +1676,7 @@ TEST_F(ParquetChunkedReaderInputLimitTest, Mixed) auto base_path = temp_env->get_temp_filepath("mixed_types"); auto test_filenames = input_limit_get_test_names(base_path); - constexpr int num_rows = 10'000'000; + constexpr int num_rows = 2'500'000; constexpr int list_size = 4; constexpr int str_size = 3; @@ -1757,16 +1755,13 @@ TEST_F(ParquetChunkedReaderInputLimitTest, Mixed) // size of the decompressed data. so 2 GB is actually not enough to hold the whole thing at // once. // - // Note that in the dictionary cases, both of these revert down to 1 chunk because the - // dictionaries dramatically shrink the size of the uncompressed data. constexpr int expected_a[] = {5, 5, 2, 1}; - input_limit_test_read(test_filenames, tbl, 0, 256 * 1024 * 1024, expected_a); - // smaller limit + input_limit_test_read(test_filenames, tbl, 0, 64 * 1024 * 1024, expected_a); constexpr int expected_b[] = {10, 9, 3, 1}; - input_limit_test_read(test_filenames, tbl, 0, 128 * 1024 * 1024, expected_b); - // include output chunking as well + input_limit_test_read(test_filenames, tbl, 0, 32 * 1024 * 1024, expected_b); + // Include output chunking as well, and verify each input format is split. constexpr int expected_c[] = {20, 18, 15, 12}; - input_limit_test_read(test_filenames, tbl, 32 * 1024 * 1024, 64 * 1024 * 1024, expected_c); + input_limit_test_read(test_filenames, tbl, 8 * 1024 * 1024, 16 * 1024 * 1024, expected_c, true); } TEST_F(ParquetChunkedReaderTest, TestChunkedReadOutOfBoundChunks) diff --git a/cpp/tests/io/parquet_writer_test.cpp b/cpp/tests/io/parquet_writer_test.cpp index cb83f6e5d22..426e5e09b04 100644 --- a/cpp/tests/io/parquet_writer_test.cpp +++ b/cpp/tests/io/parquet_writer_test.cpp @@ -2816,136 +2816,69 @@ TYPED_TEST(ParquetWriterTimestampTypeTest, TimestampsByteStreamSplit) // Base test fixture for "stress" tests struct ParquetWriterStressTest : public cudf::test::BaseFixture {}; -TEST_F(ParquetWriterStressTest, LargeTableWeakCompression) +// Keep row groups aligned with the default 5,000-row page-fragment size. +constexpr cudf::size_type stress_rows_per_row_group = 65'000; +constexpr cudf::size_type stress_num_row_groups = 4; +constexpr cudf::size_type stress_num_rows = stress_rows_per_row_group * stress_num_row_groups; + +template +void write_stress_table(std::unique_ptr const& expected) { std::vector mm_buf; mm_buf.reserve(4 * 1024 * 1024 * 16); - custom_test_memmap_sink custom_sink(&mm_buf); - - // exercises multiple rowgroups - srand(31337); - auto expected = create_random_fixed_table(16, 4 * 1024 * 1024, false); + custom_test_memmap_sink custom_sink(&mm_buf); - // write out using the custom sink (which uses device writes) + // Exercise multiple row groups without depending on the default row-group size. cudf::io::parquet_writer_options args = - cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&custom_sink}, *expected); + cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&custom_sink}, *expected) + .row_group_size_rows(stress_rows_per_row_group); cudf::io::write_parquet(args); - cudf::io::parquet_reader_options custom_args = cudf::io::parquet_reader_options::builder( - cudf::io::source_info{cudf::host_span{ - reinterpret_cast(mm_buf.data()), mm_buf.size()}}); - auto custom_tbl = cudf::io::read_parquet(custom_args); + auto const source = cudf::io::source_info{cudf::host_span{ + reinterpret_cast(mm_buf.data()), mm_buf.size()}}; + EXPECT_EQ(cudf::io::read_parquet_metadata(source).num_rowgroups(), stress_num_row_groups); + + cudf::io::parquet_reader_options custom_args = cudf::io::parquet_reader_options::builder(source); + auto custom_tbl = cudf::io::read_parquet(custom_args); CUDF_TEST_EXPECT_TABLES_EQUAL(custom_tbl.tbl->view(), expected->view()); } -TEST_F(ParquetWriterStressTest, LargeTableGoodCompression) +TEST_F(ParquetWriterStressTest, LargeTableWeakCompression) { - std::vector mm_buf; - mm_buf.reserve(4 * 1024 * 1024 * 16); - custom_test_memmap_sink custom_sink(&mm_buf); - - // exercises multiple rowgroups srand(31337); - auto expected = create_compressible_fixed_table(16, 4 * 1024 * 1024, 128 * 1024, false); - - // write out using the custom sink (which uses device writes) - cudf::io::parquet_writer_options args = - cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&custom_sink}, *expected); - cudf::io::write_parquet(args); + write_stress_table(create_random_fixed_table(16, stress_num_rows, false)); +} - cudf::io::parquet_reader_options custom_args = cudf::io::parquet_reader_options::builder( - cudf::io::source_info{cudf::host_span{ - reinterpret_cast(mm_buf.data()), mm_buf.size()}}); - auto custom_tbl = cudf::io::read_parquet(custom_args); - CUDF_TEST_EXPECT_TABLES_EQUAL(custom_tbl.tbl->view(), expected->view()); +TEST_F(ParquetWriterStressTest, LargeTableGoodCompression) +{ + srand(31337); + write_stress_table( + create_compressible_fixed_table(16, stress_num_rows, 128 * 1024, false)); } TEST_F(ParquetWriterStressTest, LargeTableWithValids) { - std::vector mm_buf; - mm_buf.reserve(4 * 1024 * 1024 * 16); - custom_test_memmap_sink custom_sink(&mm_buf); - - // exercises multiple rowgroups srand(31337); - auto expected = create_compressible_fixed_table(16, 4 * 1024 * 1024, 6, true); - - // write out using the custom sink (which uses device writes) - cudf::io::parquet_writer_options args = - cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&custom_sink}, *expected); - cudf::io::write_parquet(args); - - cudf::io::parquet_reader_options custom_args = cudf::io::parquet_reader_options::builder( - cudf::io::source_info{cudf::host_span{ - reinterpret_cast(mm_buf.data()), mm_buf.size()}}); - auto custom_tbl = cudf::io::read_parquet(custom_args); - CUDF_TEST_EXPECT_TABLES_EQUAL(custom_tbl.tbl->view(), expected->view()); + write_stress_table(create_compressible_fixed_table(16, stress_num_rows, 6, true)); } TEST_F(ParquetWriterStressTest, DeviceWriteLargeTableWeakCompression) { - std::vector mm_buf; - mm_buf.reserve(4 * 1024 * 1024 * 16); - custom_test_memmap_sink custom_sink(&mm_buf); - - // exercises multiple rowgroups srand(31337); - auto expected = create_random_fixed_table(16, 4 * 1024 * 1024, false); - - // write out using the custom sink (which uses device writes) - cudf::io::parquet_writer_options args = - cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&custom_sink}, *expected); - cudf::io::write_parquet(args); - - cudf::io::parquet_reader_options custom_args = cudf::io::parquet_reader_options::builder( - cudf::io::source_info{cudf::host_span{ - reinterpret_cast(mm_buf.data()), mm_buf.size()}}); - auto custom_tbl = cudf::io::read_parquet(custom_args); - CUDF_TEST_EXPECT_TABLES_EQUAL(custom_tbl.tbl->view(), expected->view()); + write_stress_table(create_random_fixed_table(16, stress_num_rows, false)); } TEST_F(ParquetWriterStressTest, DeviceWriteLargeTableGoodCompression) { - std::vector mm_buf; - mm_buf.reserve(4 * 1024 * 1024 * 16); - custom_test_memmap_sink custom_sink(&mm_buf); - - // exercises multiple rowgroups srand(31337); - auto expected = create_compressible_fixed_table(16, 4 * 1024 * 1024, 128 * 1024, false); - - // write out using the custom sink (which uses device writes) - cudf::io::parquet_writer_options args = - cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&custom_sink}, *expected); - cudf::io::write_parquet(args); - - cudf::io::parquet_reader_options custom_args = cudf::io::parquet_reader_options::builder( - cudf::io::source_info{cudf::host_span{ - reinterpret_cast(mm_buf.data()), mm_buf.size()}}); - auto custom_tbl = cudf::io::read_parquet(custom_args); - CUDF_TEST_EXPECT_TABLES_EQUAL(custom_tbl.tbl->view(), expected->view()); + write_stress_table( + create_compressible_fixed_table(16, stress_num_rows, 128 * 1024, false)); } TEST_F(ParquetWriterStressTest, DeviceWriteLargeTableWithValids) { - std::vector mm_buf; - mm_buf.reserve(4 * 1024 * 1024 * 16); - custom_test_memmap_sink custom_sink(&mm_buf); - - // exercises multiple rowgroups srand(31337); - auto expected = create_compressible_fixed_table(16, 4 * 1024 * 1024, 6, true); - - // write out using the custom sink (which uses device writes) - cudf::io::parquet_writer_options args = - cudf::io::parquet_writer_options::builder(cudf::io::sink_info{&custom_sink}, *expected); - cudf::io::write_parquet(args); - - cudf::io::parquet_reader_options custom_args = cudf::io::parquet_reader_options::builder( - cudf::io::source_info{cudf::host_span{ - reinterpret_cast(mm_buf.data()), mm_buf.size()}}); - auto custom_tbl = cudf::io::read_parquet(custom_args); - CUDF_TEST_EXPECT_TABLES_EQUAL(custom_tbl.tbl->view(), expected->view()); + write_stress_table(create_compressible_fixed_table(16, stress_num_rows, 6, true)); } TEST_F(ParquetWriterTest, ReturnedFooterMetadata)