Skip to content

fix(parquet): add memory control for parquet write #123

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zjw1111 merged 5 commits into from
Feb 10, 2026
Merged

fix(parquet): add memory control for parquet write #123

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
258f340
add memory control for parquet write
zjw1111 Feb 6, 2026
a45c1bf
Merge branch 'main' into fix-parquet-oom
zjw1111 Feb 6, 2026
eacebd4
fix
zjw1111 Feb 6, 2026
d850e3c
fix
zjw1111 Feb 6, 2026
d1b19a3
Merge branch 'main' into fix-parquet-oom
zjw1111 Feb 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion src/paimon/format/parquet/file_reader_wrapper_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
#include "paimon/common/utils/arrow/status_utils.h"
#include "paimon/common/utils/path_util.h"
#include "paimon/format/parquet/parquet_field_id_converter.h"
#include "paimon/format/parquet/parquet_format_defs.h"
#include "paimon/format/parquet/parquet_format_writer.h"
#include "paimon/format/parquet/parquet_input_stream_impl.h"
#include "paimon/fs/file_system.h"
Expand Down Expand Up @@ -150,7 +151,8 @@ class FileReaderWrapperTest : public ::testing::Test {
auto writer_properties = builder.build();
ASSERT_OK_AND_ASSIGN(
std::shared_ptr<ParquetFormatWriter> format_writer,
ParquetFormatWriter::Create(out, arrow_schema, writer_properties, arrow_pool_));
ParquetFormatWriter::Create(out, arrow_schema, writer_properties,
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, arrow_pool_));

AddRecordBatchOnce(format_writer, struct_type, /*record_batch_size=*/row_count,
/*offset=*/0);
Expand Down
4 changes: 2 additions & 2 deletions src/paimon/format/parquet/parquet_file_batch_reader.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -265,8 +265,8 @@ Result<::parquet::ArrowReaderProperties> ParquetFileBatchReader::CreateArrowRead
const std::shared_ptr<arrow::MemoryPool>& pool,
const std::map<std::string, std::string>& options, int32_t batch_size) {
PAIMON_ASSIGN_OR_RAISE(bool use_threads,
OptionsUtils::GetValueFromMap<bool>(options, PARQUET_READ_USE_THREADS,
DEFAULT_PARQUET_READ_USE_THREADS));
OptionsUtils::GetValueFromMap<bool>(options, PARQUET_USE_MULTI_THREAD,
DEFAULT_PARQUET_USE_MULTI_THREAD));

::parquet::ArrowReaderProperties arrow_reader_props;
// TODO(jinli.zjw): set more ArrowReaderProperties (compare with java)
Expand Down
8 changes: 4 additions & 4 deletions src/paimon/format/parquet/parquet_file_batch_reader_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@
#include "arrow/io/caching.h"
#include "arrow/io/interfaces.h"
#include "arrow/ipc/json_simple.h"
#include "arrow/util/thread_pool.h"
#include "gtest/gtest.h"
#include "paimon/common/types/data_field.h"
#include "paimon/common/utils/arrow/mem_utils.h"
Expand Down Expand Up @@ -112,7 +111,8 @@ class ParquetFileBatchReaderTest : public ::testing::Test,
enable_dictionary ? builder.enable_dictionary() : builder.disable_dictionary();
auto writer_properties = builder.build();
ASSERT_OK_AND_ASSIGN(auto format_writer, ParquetFormatWriter::Create(
out, arrow_schema, writer_properties, pool_));
out, arrow_schema, writer_properties,
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, pool_));

auto arrow_array = std::make_unique<ArrowArray>();
ASSERT_TRUE(arrow::ExportArray(*src_array, arrow_array.get()).ok());
Expand Down Expand Up @@ -398,7 +398,7 @@ TEST_F(ParquetFileBatchReaderTest, TestCreateArrowReaderProperties) {
ASSERT_EQ(arrow_reader_properties.cache_options(), arrow::io::CacheOptions::Defaults());
}
{
std::map<std::string, std::string> options = {{PARQUET_READ_USE_THREADS, "false"}};
std::map<std::string, std::string> options = {{PARQUET_USE_MULTI_THREAD, "false"}};
int32_t batch_size = 1024;
ASSERT_OK_AND_ASSIGN(
auto arrow_reader_properties,
Expand All @@ -408,7 +408,7 @@ TEST_F(ParquetFileBatchReaderTest, TestCreateArrowReaderProperties) {
{
int original_capacity = GetArrowCpuThreadPoolCapacity();
ASSERT_OK(SetArrowCpuThreadPoolCapacity(6));
std::map<std::string, std::string> options = {{PARQUET_READ_USE_THREADS, "true"}};
std::map<std::string, std::string> options = {{PARQUET_USE_MULTI_THREAD, "true"}};
int32_t batch_size = 1024;
ASSERT_OK_AND_ASSIGN(
auto arrow_reader_properties,
Expand Down
8 changes: 5 additions & 3 deletions src/paimon/format/parquet/parquet_format_defs.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ namespace paimon::parquet {
static inline const char PARQUET_BLOCK_SIZE[] = "parquet.block.size";
static inline const char PARQUET_PAGE_SIZE[] = "parquet.page.size";
static inline const char PARQUET_DICTIONARY_PAGE_SIZE[] = "parquet.dictionary.page.size";
static inline const char PARQUET_ENABLE_DICTIONARY[] = "parquet.enable.dictionary";
static inline const char PARQUET_ENABLE_DICTIONARY[] = "parquet.enable-dictionary";
static inline const char PARQUET_WRITER_VERSION[] = "parquet.writer.version";
static inline const char PARQUET_WRITE_MAX_ROW_GROUP_LENGTH[] =
"parquet.write.max-row-group-length";
Expand All @@ -34,10 +34,12 @@ static inline const char PARQUET_COMPRESSION_CODEC_ZSTD_LEVEL[] =
"parquet.compression.codec.zstd.level";
static inline const char PARQUET_COMPRESSION_CODEC_ZLIB_LEVEL[] = "zlib.compress.level";
static inline const char PARQUET_COMPRESSION_CODEC_BROTLI_LEVEL[] = "compression.brotli.quality";
static inline const char PARQUET_WRITER_MAX_MEMORY_USE[] = "parquet.writer.max.memory.use";
static constexpr uint64_t DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE = 512 * 1024 * 1024; // 512MB

// read
static inline const char PARQUET_READ_USE_THREADS[] = "parquet.read.use-threads";
static inline const bool DEFAULT_PARQUET_READ_USE_THREADS = true;
static inline const char PARQUET_USE_MULTI_THREAD[] = "parquet.use-multi-thread";
static inline const bool DEFAULT_PARQUET_USE_MULTI_THREAD = true;
static inline const char PARQUET_READ_CACHE_OPTION_LAZY[] = "parquet.read.cache-option.lazy";
static inline const char PARQUET_READ_CACHE_OPTION_PREFETCH_LIMIT[] =
"parquet.read.cache-option.prefetch-limit";
Expand Down
11 changes: 8 additions & 3 deletions src/paimon/format/parquet/parquet_format_writer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ namespace paimon::parquet {
Result<std::unique_ptr<ParquetFormatWriter>> ParquetFormatWriter::Create(
const std::shared_ptr<OutputStream>& output_stream,
const std::shared_ptr<arrow::Schema>& schema,
const std::shared_ptr<::parquet::WriterProperties>& writer_properties,
const std::shared_ptr<::parquet::WriterProperties>& writer_properties, uint64_t max_memory_use,
const std::shared_ptr<arrow::MemoryPool>& pool) {
auto out = std::make_shared<ParquetOutputStreamImpl>(output_stream);
::parquet::ArrowWriterProperties::Builder arrow_properties_builder;
Expand All @@ -54,12 +54,15 @@ Result<std::unique_ptr<ParquetFormatWriter>> ParquetFormatWriter::Create(
::parquet::arrow::FileWriter::Open(*schema, pool.get(), out, writer_properties,
arrow_writer_properties));
return std::unique_ptr<ParquetFormatWriter>(
new ParquetFormatWriter(std::move(file_writer), out, schema, pool));
new ParquetFormatWriter(std::move(file_writer), out, schema, max_memory_use, pool));
}

Status ParquetFormatWriter::AddBatch(ArrowArray* batch) {
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<::arrow::RecordBatch> record_batch,
arrow::ImportRecordBatch(batch, schema_));
if (static_cast<uint64_t>(pool_->bytes_allocated()) > max_memory_use_) {
PAIMON_RETURN_NOT_OK_FROM_ARROW(writer_->NewBufferedRowGroup());
}
Comment thread
zjw1111 marked this conversation as resolved.
Comment thread
zjw1111 marked this conversation as resolved.
PAIMON_RETURN_NOT_OK_FROM_ARROW(writer_->WriteRecordBatch(*record_batch));
total_records_written_ += (*record_batch).num_rows();
return Status::OK();
Expand Down Expand Up @@ -92,11 +95,13 @@ Result<uint64_t> ParquetFormatWriter::GetEstimateLength() const {
ParquetFormatWriter::ParquetFormatWriter(std::unique_ptr<::parquet::arrow::FileWriter> writer,
const std::shared_ptr<ParquetOutputStreamImpl>& out,
const std::shared_ptr<arrow::Schema>& schema,
uint64_t max_memory_use,
const std::shared_ptr<arrow::MemoryPool>& pool)
: pool_(pool),
out_(out),
writer_(std::move(writer)),
schema_(schema),
metrics_(std::make_shared<MetricsImpl>()) {}
metrics_(std::make_shared<MetricsImpl>()),
max_memory_use_(max_memory_use) {}

} // namespace paimon::parquet
5 changes: 3 additions & 2 deletions src/paimon/format/parquet/parquet_format_writer.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ class ParquetFormatWriter : public FormatWriter {
const std::shared_ptr<OutputStream>& output_stream,
const std::shared_ptr<arrow::Schema>& schema,
const std::shared_ptr<::parquet::WriterProperties>& writer_properties,
const std::shared_ptr<arrow::MemoryPool>& pool);
uint64_t max_memory_use, const std::shared_ptr<arrow::MemoryPool>& pool);

Status AddBatch(ArrowArray* batch) override;

Expand All @@ -71,7 +71,7 @@ class ParquetFormatWriter : public FormatWriter {
private:
ParquetFormatWriter(std::unique_ptr<::parquet::arrow::FileWriter> writer,
const std::shared_ptr<ParquetOutputStreamImpl>& out,
const std::shared_ptr<arrow::Schema>& schema,
const std::shared_ptr<arrow::Schema>& schema, uint64_t max_memory_use,
const std::shared_ptr<arrow::MemoryPool>& pool);

Result<uint64_t> GetEstimateLength() const;
Expand All @@ -82,6 +82,7 @@ class ParquetFormatWriter : public FormatWriter {
std::shared_ptr<arrow::Schema> schema_;
std::shared_ptr<Metrics> metrics_;
int64_t total_records_written_ = 0;
uint64_t max_memory_use_;
};

} // namespace paimon::parquet
Loading
Loading