Skip to content

Commit

Permalink
Fix batch processing for parquet writer (#13438)
Browse files Browse the repository at this point in the history
In parquet writer, the input table is divided into multiple batches (at 1GB limit), each batch is processed and flushed to sink one after another. The buffers storing data for processing each batch are reused among batches. This is to reduce peak GPU memory usage.

Unfortunately, in order to support retry mechanism, we have to have separate buffers for each batch. This is equivalent to always having one batch. The benefit of batch processing is stripped away.  In #13076, we expect to keep data for all batches but failed to do that, causing a bug reported in #13414.

This PR fixes the issue introduced in #13076. And since we have to strip away the benefit of batch processing, peak memory usage may go up.

Flag this as `breaking` because peak GPU memory usage may go up and cause the downstream application to crash.

Note that this PR is a temporary fix for the outstanding issue. With this fix, the batch processing mechanism no longer gives any benefit for reducing peak memory usage. We consider removing all the batch processing code completely in the follow-up work, which involves a lot more changes.

Closes #13414.

Authors:
  - Nghia Truong (https://github.com/ttnghia)

Approvers:
  - Vukasin Milovanovic (https://github.com/vuule)
  - Lawrence Mitchell (https://github.com/wence-)
  - Benjamin Zaitlen (https://github.com/quasiben)

URL: #13438
  • Loading branch information
ttnghia authored May 25, 2023
1 parent 960cc42 commit ae375d2
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion cpp/src/io/parquet/writer_impl.cu
Original file line number Diff line number Diff line change
Expand Up @@ -1771,10 +1771,15 @@ auto convert_table_to_parquet_data(table_input_metadata& table_meta,
std::fill_n(std::back_inserter(rg_to_part), num_rg_in_part[p], p);
}

// Batch processing is no longer supported.
// This line disables batch processing (so batch size will no longer be limited at 1GB as before).
// TODO: All the relevant code will be removed in the follow-up work:
// https://github.com/rapidsai/cudf/issues/13440
auto const max_bytes_in_batch = std::numeric_limits<size_t>::max();

// Initialize batches of rowgroups to encode (mainly to limit peak memory usage)
std::vector<size_type> batch_list;
size_type num_pages = 0;
size_t max_bytes_in_batch = 1024 * 1024 * 1024; // 1GB - TODO: Tune this
size_t max_uncomp_bfr_size = 0;
size_t max_comp_bfr_size = 0;
size_t max_chunk_bfr_size = 0;
Expand Down

0 comments on commit ae375d2

Please sign in to comment.