Skip to content

Commit 623db38

Browse files
zjw1111zjw1111
authored
fix(parquet): add memory control for parquet write (#123)
1 parent 2016e9f commit 623db38

11 files changed

Lines changed: 166 additions & 43 deletions

‎src/paimon/format/parquet/file_reader_wrapper_test.cpp‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
#include "paimon/common/utils/arrow/status_utils.h"
3333
#include "paimon/common/utils/path_util.h"
3434
#include "paimon/format/parquet/parquet_field_id_converter.h"
35+
#include "paimon/format/parquet/parquet_format_defs.h"
3536
#include "paimon/format/parquet/parquet_format_writer.h"
3637
#include "paimon/format/parquet/parquet_input_stream_impl.h"
3738
#include "paimon/fs/file_system.h"
@@ -150,7 +151,8 @@ class FileReaderWrapperTest : public ::testing::Test {
150151
auto writer_properties = builder.build();
151152
ASSERT_OK_AND_ASSIGN(
152153
std::shared_ptr<ParquetFormatWriter> format_writer,
153-
ParquetFormatWriter::Create(out, arrow_schema, writer_properties, arrow_pool_));
154+
ParquetFormatWriter::Create(out, arrow_schema, writer_properties,
155+
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, arrow_pool_));
154156

155157
AddRecordBatchOnce(format_writer, struct_type, /*record_batch_size=*/row_count,
156158
/*offset=*/0);

‎src/paimon/format/parquet/parquet_file_batch_reader.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -265,8 +265,8 @@ Result<::parquet::ArrowReaderProperties> ParquetFileBatchReader::CreateArrowRead
265265
const std::shared_ptr<arrow::MemoryPool>& pool,
266266
const std::map<std::string, std::string>& options, int32_t batch_size) {
267267
PAIMON_ASSIGN_OR_RAISE(bool use_threads,
268-
OptionsUtils::GetValueFromMap<bool>(options, PARQUET_READ_USE_THREADS,
269-
DEFAULT_PARQUET_READ_USE_THREADS));
268+
OptionsUtils::GetValueFromMap<bool>(options, PARQUET_USE_MULTI_THREAD,
269+
DEFAULT_PARQUET_USE_MULTI_THREAD));
270270

271271
::parquet::ArrowReaderProperties arrow_reader_props;
272272
// TODO(jinli.zjw): set more ArrowReaderProperties (compare with java)

‎src/paimon/format/parquet/parquet_file_batch_reader_test.cpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@
2929
#include "arrow/io/caching.h"
3030
#include "arrow/io/interfaces.h"
3131
#include "arrow/ipc/json_simple.h"
32-
#include "arrow/util/thread_pool.h"
3332
#include "gtest/gtest.h"
3433
#include "paimon/common/types/data_field.h"
3534
#include "paimon/common/utils/arrow/mem_utils.h"
@@ -112,7 +111,8 @@ class ParquetFileBatchReaderTest : public ::testing::Test,
112111
enable_dictionary ? builder.enable_dictionary() : builder.disable_dictionary();
113112
auto writer_properties = builder.build();
114113
ASSERT_OK_AND_ASSIGN(auto format_writer, ParquetFormatWriter::Create(
115-
out, arrow_schema, writer_properties, pool_));
114+
out, arrow_schema, writer_properties,
115+
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, pool_));
116116

117117
auto arrow_array = std::make_unique<ArrowArray>();
118118
ASSERT_TRUE(arrow::ExportArray(*src_array, arrow_array.get()).ok());
@@ -398,7 +398,7 @@ TEST_F(ParquetFileBatchReaderTest, TestCreateArrowReaderProperties) {
398398
ASSERT_EQ(arrow_reader_properties.cache_options(), arrow::io::CacheOptions::Defaults());
399399
}
400400
{
401-
std::map<std::string, std::string> options = {{PARQUET_READ_USE_THREADS, "false"}};
401+
std::map<std::string, std::string> options = {{PARQUET_USE_MULTI_THREAD, "false"}};
402402
int32_t batch_size = 1024;
403403
ASSERT_OK_AND_ASSIGN(
404404
auto arrow_reader_properties,
@@ -408,7 +408,7 @@ TEST_F(ParquetFileBatchReaderTest, TestCreateArrowReaderProperties) {
408408
{
409409
int original_capacity = GetArrowCpuThreadPoolCapacity();
410410
ASSERT_OK(SetArrowCpuThreadPoolCapacity(6));
411-
std::map<std::string, std::string> options = {{PARQUET_READ_USE_THREADS, "true"}};
411+
std::map<std::string, std::string> options = {{PARQUET_USE_MULTI_THREAD, "true"}};
412412
int32_t batch_size = 1024;
413413
ASSERT_OK_AND_ASSIGN(
414414
auto arrow_reader_properties,

‎src/paimon/format/parquet/parquet_format_defs.h‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ namespace paimon::parquet {
2424
static inline const char PARQUET_BLOCK_SIZE[] = "parquet.block.size";
2525
static inline const char PARQUET_PAGE_SIZE[] = "parquet.page.size";
2626
static inline const char PARQUET_DICTIONARY_PAGE_SIZE[] = "parquet.dictionary.page.size";
27-
static inline const char PARQUET_ENABLE_DICTIONARY[] = "parquet.enable.dictionary";
27+
static inline const char PARQUET_ENABLE_DICTIONARY[] = "parquet.enable-dictionary";
2828
static inline const char PARQUET_WRITER_VERSION[] = "parquet.writer.version";
2929
static inline const char PARQUET_WRITE_MAX_ROW_GROUP_LENGTH[] =
3030
"parquet.write.max-row-group-length";
@@ -34,10 +34,12 @@ static inline const char PARQUET_COMPRESSION_CODEC_ZSTD_LEVEL[] =
3434
"parquet.compression.codec.zstd.level";
3535
static inline const char PARQUET_COMPRESSION_CODEC_ZLIB_LEVEL[] = "zlib.compress.level";
3636
static inline const char PARQUET_COMPRESSION_CODEC_BROTLI_LEVEL[] = "compression.brotli.quality";
37+
static inline const char PARQUET_WRITER_MAX_MEMORY_USE[] = "parquet.writer.max.memory.use";
38+
static constexpr uint64_t DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE = 512 * 1024 * 1024; // 512MB
3739

3840
// read
39-
static inline const char PARQUET_READ_USE_THREADS[] = "parquet.read.use-threads";
40-
static inline const bool DEFAULT_PARQUET_READ_USE_THREADS = true;
41+
static inline const char PARQUET_USE_MULTI_THREAD[] = "parquet.use-multi-thread";
42+
static inline const bool DEFAULT_PARQUET_USE_MULTI_THREAD = true;
4143
static inline const char PARQUET_READ_CACHE_OPTION_LAZY[] = "parquet.read.cache-option.lazy";
4244
static inline const char PARQUET_READ_CACHE_OPTION_PREFETCH_LIMIT[] =
4345
"parquet.read.cache-option.prefetch-limit";

‎src/paimon/format/parquet/parquet_format_writer.cpp‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ namespace paimon::parquet {
4343
Result<std::unique_ptr<ParquetFormatWriter>> ParquetFormatWriter::Create(
4444
const std::shared_ptr<OutputStream>& output_stream,
4545
const std::shared_ptr<arrow::Schema>& schema,
46-
const std::shared_ptr<::parquet::WriterProperties>& writer_properties,
46+
const std::shared_ptr<::parquet::WriterProperties>& writer_properties, uint64_t max_memory_use,
4747
const std::shared_ptr<arrow::MemoryPool>& pool) {
4848
auto out = std::make_shared<ParquetOutputStreamImpl>(output_stream);
4949
::parquet::ArrowWriterProperties::Builder arrow_properties_builder;
@@ -54,12 +54,15 @@ Result<std::unique_ptr<ParquetFormatWriter>> ParquetFormatWriter::Create(
5454
::parquet::arrow::FileWriter::Open(*schema, pool.get(), out, writer_properties,
5555
arrow_writer_properties));
5656
return std::unique_ptr<ParquetFormatWriter>(
57-
new ParquetFormatWriter(std::move(file_writer), out, schema, pool));
57+
new ParquetFormatWriter(std::move(file_writer), out, schema, max_memory_use, pool));
5858
}
5959

6060
Status ParquetFormatWriter::AddBatch(ArrowArray* batch) {
6161
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<::arrow::RecordBatch> record_batch,
6262
arrow::ImportRecordBatch(batch, schema_));
63+
if (static_cast<uint64_t>(pool_->bytes_allocated()) > max_memory_use_) {
64+
PAIMON_RETURN_NOT_OK_FROM_ARROW(writer_->NewBufferedRowGroup());
65+
}
6366
PAIMON_RETURN_NOT_OK_FROM_ARROW(writer_->WriteRecordBatch(*record_batch));
6467
total_records_written_ += (*record_batch).num_rows();
6568
return Status::OK();
@@ -92,11 +95,13 @@ Result<uint64_t> ParquetFormatWriter::GetEstimateLength() const {
9295
ParquetFormatWriter::ParquetFormatWriter(std::unique_ptr<::parquet::arrow::FileWriter> writer,
9396
const std::shared_ptr<ParquetOutputStreamImpl>& out,
9497
const std::shared_ptr<arrow::Schema>& schema,
98+
uint64_t max_memory_use,
9599
const std::shared_ptr<arrow::MemoryPool>& pool)
96100
: pool_(pool),
97101
out_(out),
98102
writer_(std::move(writer)),
99103
schema_(schema),
100-
metrics_(std::make_shared<MetricsImpl>()) {}
104+
metrics_(std::make_shared<MetricsImpl>()),
105+
max_memory_use_(max_memory_use) {}
101106

102107
} // namespace paimon::parquet

‎src/paimon/format/parquet/parquet_format_writer.h‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ class ParquetFormatWriter : public FormatWriter {
5454
const std::shared_ptr<OutputStream>& output_stream,
5555
const std::shared_ptr<arrow::Schema>& schema,
5656
const std::shared_ptr<::parquet::WriterProperties>& writer_properties,
57-
const std::shared_ptr<arrow::MemoryPool>& pool);
57+
uint64_t max_memory_use, const std::shared_ptr<arrow::MemoryPool>& pool);
5858

5959
Status AddBatch(ArrowArray* batch) override;
6060

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

7777
Result<uint64_t> GetEstimateLength() const;
@@ -82,6 +82,7 @@ class ParquetFormatWriter : public FormatWriter {
8282
std::shared_ptr<arrow::Schema> schema_;
8383
std::shared_ptr<Metrics> metrics_;
8484
int64_t total_records_written_ = 0;
85+
uint64_t max_memory_use_;
8586
};
8687

8788
} // namespace paimon::parquet

0 commit comments

Comments
 (0)