fix(parquet): add memory control for parquet write#123
Merged
Conversation
zjw1111
changed the title
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds memory control to the Parquet writer to prevent out-of-memory errors during write operations. Previously, all data batches in a row group were held in memory until the entire row group was flushed, which could lead to excessive memory usage. The new implementation monitors memory usage and forces a row group flush when memory exceeds a configurable threshold (default 512MB).
Changes:
- Added memory control logic that checks pool memory usage before writing each batch and forces row group flush when threshold is exceeded
- Added new configuration option
parquet.writer.max.memory.usewith a default value of 512MB - Updated all ParquetFormatWriter::Create calls to include the new max_memory_use parameter
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| src/paimon/format/parquet/parquet_format_defs.h | Added PARQUET_WRITER_MAX_MEMORY_USE constant and DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE (512MB) |
| src/paimon/format/parquet/parquet_format_writer.h | Added max_memory_use parameter to Create method and constructor, added max_memory_use_ member variable |
| src/paimon/format/parquet/parquet_format_writer.cpp | Implemented memory check in AddBatch that triggers NewBufferedRowGroup when threshold exceeded |
| src/paimon/format/parquet/parquet_writer_builder.h | Removed unused includes (arrow/util/compression.h, paimon/fs/file_system_factory.h) |
| src/paimon/format/parquet/parquet_writer_builder.cpp | Added logic to extract max_memory_use from options and pass to ParquetFormatWriter::Create |
| src/paimon/format/parquet/parquet_format_writer_test.cpp | Added TestMemoryControl test case, updated all Create calls with new parameter, added test infrastructure for null values |
| src/paimon/format/parquet/predicate_pushdown_test.cpp | Updated ParquetFormatWriter::Create call with DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE parameter |
| src/paimon/format/parquet/parquet_stats_extractor_test.cpp | Added include for parquet_format_defs.h, updated Create calls with new parameter |
| src/paimon/format/parquet/parquet_file_batch_reader_test.cpp | Removed unused include (arrow/util/thread_pool.h), updated Create call with new parameter |
| src/paimon/format/parquet/file_reader_wrapper_test.cpp | Added include for parquet_format_defs.h, updated Create call with new parameter |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Linked issue: close #xxx
In parquet writer, the memory occupied by all batches of data is only released after the entire row group has been flushed. This can lead to excessive memory usage during the writing process, posing a risk of out-of-memory (OOM) errors. This PR adds memory control, which forces flush the row group when the pool's memory usage reaches a specified threshold.
Tests
ParquetFormatWriterTest.TestMemoryControl
API and Format
Documentation