Skip to content

Commit 258f340

Browse files
zjw1111zjw1111
committed
add memory control for parquet write
1 parent 4291b6e commit 258f340

10 files changed

Lines changed: 114 additions & 31 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, arrow_pool_,
155+
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE));
154156

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

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

Lines changed: 2 additions & 2 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, pool_,
115+
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE));
116116

117117
auto arrow_array = std::make_unique<ArrowArray>();
118118
ASSERT_TRUE(arrow::ExportArray(*src_array, arrow_array.get()).ok());

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ 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
3941
static inline const char PARQUET_READ_USE_THREADS[] = "parquet.read.use-threads";

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

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ Result<std::unique_ptr<ParquetFormatWriter>> ParquetFormatWriter::Create(
4444
const std::shared_ptr<OutputStream>& output_stream,
4545
const std::shared_ptr<arrow::Schema>& schema,
4646
const std::shared_ptr<::parquet::WriterProperties>& writer_properties,
47-
const std::shared_ptr<arrow::MemoryPool>& pool) {
47+
const std::shared_ptr<arrow::MemoryPool>& pool, uint64_t max_memory_use) {
4848
auto out = std::make_shared<ParquetOutputStreamImpl>(output_stream);
4949
::parquet::ArrowWriterProperties::Builder arrow_properties_builder;
5050
auto arrow_writer_properties =
@@ -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, pool, max_memory_use));
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,
95-
const std::shared_ptr<arrow::MemoryPool>& pool)
98+
const std::shared_ptr<arrow::MemoryPool>& pool,
99+
uint64_t max_memory_use)
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+
const std::shared_ptr<arrow::MemoryPool>& pool, uint64_t max_memory_use);
5858

5959
Status AddBatch(ArrowArray* batch) override;
6060

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

7777
Result<uint64_t> GetEstimateLength() const;
7878

@@ -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

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

Lines changed: 83 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@
3636
#include "paimon/common/utils/arrow/mem_utils.h"
3737
#include "paimon/common/utils/date_time_utils.h"
3838
#include "paimon/common/utils/path_util.h"
39+
#include "paimon/format/file_format.h"
40+
#include "paimon/format/file_format_factory.h"
3941
#include "paimon/format/parquet/parquet_field_id_converter.h"
4042
#include "paimon/format/parquet/parquet_format_defs.h"
4143
#include "paimon/fs/file_system.h"
@@ -84,8 +86,8 @@ class ParquetFormatWriterTest : public ::testing::Test {
8486
}
8587

8688
std::shared_ptr<arrow::Array> PrepareArray(const std::shared_ptr<arrow::DataType>& data_type,
87-
int32_t record_batch_size,
88-
int32_t offset = 0) const {
89+
int32_t record_batch_size, int32_t offset = 0,
90+
bool all_null_value = false) const {
8991
arrow::StructBuilder struct_builder(
9092
data_type, arrow::default_memory_pool(),
9193
{std::make_shared<arrow::StringBuilder>(), std::make_shared<arrow::Int32Builder>(),
@@ -95,24 +97,31 @@ class ParquetFormatWriterTest : public ::testing::Test {
9597
auto bool_builder = static_cast<arrow::BooleanBuilder*>(struct_builder.field_builder(2));
9698
for (int32_t i = 0 + offset; i < record_batch_size + offset; ++i) {
9799
EXPECT_TRUE(struct_builder.Append().ok());
98-
EXPECT_TRUE(string_builder->Append("str_" + std::to_string(i)).ok());
99-
if (i % 3 == 0) {
100-
// test null
100+
if (all_null_value) {
101+
EXPECT_TRUE(string_builder->AppendNull().ok());
101102
EXPECT_TRUE(int_builder->AppendNull().ok());
103+
EXPECT_TRUE(bool_builder->AppendNull().ok());
102104
} else {
103-
EXPECT_TRUE(int_builder->Append(i).ok());
105+
EXPECT_TRUE(string_builder->Append("str_" + std::to_string(i)).ok());
106+
if (i % 3 == 0) {
107+
// test null
108+
EXPECT_TRUE(int_builder->AppendNull().ok());
109+
} else {
110+
EXPECT_TRUE(int_builder->Append(i).ok());
111+
}
112+
EXPECT_TRUE(bool_builder->Append(static_cast<bool>(i % 2)).ok());
104113
}
105-
EXPECT_TRUE(bool_builder->Append(static_cast<bool>(i % 2)).ok());
106114
}
107115
std::shared_ptr<arrow::Array> array;
108116
EXPECT_TRUE(struct_builder.Finish(&array).ok());
109117
return array;
110118
}
111119

112-
void AddRecordBatchOnce(const std::shared_ptr<ParquetFormatWriter>& format_writer,
120+
void AddRecordBatchOnce(const std::shared_ptr<FormatWriter>& format_writer,
113121
const std::shared_ptr<arrow::DataType>& struct_type,
114-
int32_t record_batch_size, int32_t offset) const {
115-
auto array = PrepareArray(struct_type, record_batch_size, offset);
122+
int32_t record_batch_size, int32_t offset,
123+
bool all_null_value = false) const {
124+
auto array = PrepareArray(struct_type, record_batch_size, offset, all_null_value);
116125
auto arrow_array = std::make_unique<ArrowArray>();
117126
ASSERT_TRUE(arrow::ExportArray(*array, arrow_array.get()).ok());
118127
auto batch = std::make_shared<RecordBatch>(
@@ -196,7 +205,8 @@ TEST_F(ParquetFormatWriterTest, TestWriteWithVariousBatchSize) {
196205
auto writer_properties = builder.build();
197206
ASSERT_OK_AND_ASSIGN(
198207
auto format_writer,
199-
ParquetFormatWriter::Create(out, arrow_schema, writer_properties, arrow_pool_));
208+
ParquetFormatWriter::Create(out, arrow_schema, writer_properties, arrow_pool_,
209+
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE));
200210
auto array = PrepareArray(struct_type, record_batch_size);
201211
auto arrow_array = std::make_unique<ArrowArray>();
202212
ASSERT_TRUE(arrow::ExportArray(*array, arrow_array.get()).ok());
@@ -229,7 +239,8 @@ TEST_F(ParquetFormatWriterTest, TestWriteMultipleTimes) {
229239
auto writer_properties = builder.build();
230240
ASSERT_OK_AND_ASSIGN(
231241
std::shared_ptr<ParquetFormatWriter> format_writer,
232-
ParquetFormatWriter::Create(out, arrow_schema, writer_properties, arrow_pool_));
242+
ParquetFormatWriter::Create(out, arrow_schema, writer_properties, arrow_pool_,
243+
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE));
233244

234245
// add batch first time, 6 rows
235246
AddRecordBatchOnce(format_writer, struct_type, 6, 0);
@@ -271,7 +282,8 @@ TEST_F(ParquetFormatWriterTest, TestGetEstimateLength) {
271282
auto writer_properties = builder.build();
272283
ASSERT_OK_AND_ASSIGN(
273284
std::shared_ptr<ParquetFormatWriter> format_writer,
274-
ParquetFormatWriter::Create(out, arrow_schema, writer_properties, arrow_pool_));
285+
ParquetFormatWriter::Create(out, arrow_schema, writer_properties, arrow_pool_,
286+
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE));
275287

276288
// add batch first time, 1 row
277289
AddRecordBatchOnce(format_writer, struct_type, 1, 0);
@@ -289,6 +301,59 @@ TEST_F(ParquetFormatWriterTest, TestGetEstimateLength) {
289301
ASSERT_TRUE(format_writer->Finish().ok());
290302
}
291303

304+
TEST_F(ParquetFormatWriterTest, TestMemoryControl) {
305+
auto check_result = [&](bool all_null_value, uint64_t max_memory_use) {
306+
ASSERT_OK_AND_ASSIGN(
307+
std::unique_ptr<FileFormat> file_format,
308+
FileFormatFactory::Get(
309+
"parquet", {{Options::FILE_FORMAT, "parquet"},
310+
{Options::MANIFEST_FORMAT, "parquet"},
311+
{"parquet.writer.max.memory.use", std::to_string(max_memory_use)}}));
312+
313+
std::shared_ptr<MemoryPool> pool = GetMemoryPool();
314+
auto schema_pair = PrepareArrowSchema();
315+
const auto& arrow_schema = schema_pair.first;
316+
const auto& struct_type = schema_pair.second;
317+
int32_t batch_size = 4096;
318+
319+
auto c_schema = std::make_unique<::ArrowSchema>();
320+
ASSERT_TRUE(arrow::ExportSchema(*arrow_schema, c_schema.get()).ok());
321+
ASSERT_OK_AND_ASSIGN(auto writer_builder,
322+
file_format->CreateWriterBuilder(c_schema.get(), batch_size));
323+
ASSERT_OK_AND_ASSIGN(
324+
std::shared_ptr<OutputStream> out,
325+
fs_->Create(
326+
PathUtil::JoinPath(dir_->Str(), std::to_string(all_null_value) +
327+
std::to_string(max_memory_use) + ".parquet"),
328+
/*overwrite=*/false));
329+
ASSERT_OK_AND_ASSIGN(std::shared_ptr<FormatWriter> writer,
330+
writer_builder->WithMemoryPool(pool)->Build(out, "uncompressed"));
331+
332+
auto array = PrepareArray(struct_type, batch_size, /*offset=*/0, all_null_value);
333+
for (int32_t i = 0; i < 2000; ++i) {
334+
auto arrow_array = std::make_unique<ArrowArray>();
335+
ASSERT_TRUE(arrow::ExportArray(*array, arrow_array.get()).ok());
336+
auto batch = std::make_shared<RecordBatch>(
337+
/*partition=*/std::map<std::string, std::string>(), /*bucket=*/-1,
338+
/*row_kinds=*/std::vector<RecordBatch::RowKind>(), arrow_array.get());
339+
ASSERT_OK(writer->AddBatch(batch->GetData()));
340+
ASSERT_OK(writer->Flush());
341+
}
342+
343+
ASSERT_OK(writer->Flush());
344+
ASSERT_OK(writer->Finish());
345+
ASSERT_OK(out->Flush());
346+
ASSERT_OK(out->Close());
347+
uint64_t actual_max_mem = pool->MaxMemoryUsage();
348+
ASSERT_GT(actual_max_mem, max_memory_use);
349+
ASSERT_LT(actual_max_mem, max_memory_use * 1.5); // allow 50% overhead
350+
};
351+
check_result(/*all_null_value=*/true, /*max_memory_use=*/20 * 1024 * 1024); // 20MB
352+
check_result(/*all_null_value=*/true, /*max_memory_use=*/40 * 1024 * 1024); // 40MB
353+
check_result(/*all_null_value=*/false, /*max_memory_use=*/20 * 1024 * 1024); // 20MB
354+
check_result(/*all_null_value=*/false, /*max_memory_use=*/40 * 1024 * 1024); // 40MB
355+
}
356+
292357
TEST_F(ParquetFormatWriterTest, TestTimestampType) {
293358
auto timezone = DateTimeUtils::GetLocalTimezoneName();
294359
arrow::FieldVector fields = {
@@ -298,15 +363,17 @@ TEST_F(ParquetFormatWriterTest, TestTimestampType) {
298363
arrow::field("ts_nano", arrow::timestamp(arrow::TimeUnit::NANO)),
299364
arrow::field("ts_utc1", arrow::timestamp(arrow::TimeUnit::SECOND, timezone)),
300365
arrow::field("ts_utc2", arrow::timestamp(arrow::TimeUnit::MICRO, timezone))};
366+
[[maybe_unused]] auto test = new arrow::Schema(fields); // for test
301367

302368
std::string file_path = PathUtil::JoinPath(dir_->Str(), "timezone.parquet");
303369
ASSERT_OK_AND_ASSIGN(std::shared_ptr<OutputStream> out,
304370
fs_->Create(file_path, /*overwrite=*/true));
305371
::parquet::WriterProperties::Builder builder;
306372
auto writer_properties = builder.build();
307-
ASSERT_OK_AND_ASSIGN(std::shared_ptr<ParquetFormatWriter> format_writer,
308-
ParquetFormatWriter::Create(out, std::make_shared<arrow::Schema>(fields),
309-
writer_properties, arrow_pool_));
373+
ASSERT_OK_AND_ASSIGN(
374+
std::shared_ptr<ParquetFormatWriter> format_writer,
375+
ParquetFormatWriter::Create(out, std::make_shared<arrow::Schema>(fields), writer_properties,
376+
arrow_pool_, DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE));
310377

311378
auto array = std::dynamic_pointer_cast<arrow::StructArray>(
312379
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(fields), R"([

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
#include "paimon/core/stats/simple_stats.h"
4040
#include "paimon/core/stats/simple_stats_converter.h"
4141
#include "paimon/format/column_stats.h"
42+
#include "paimon/format/parquet/parquet_format_defs.h"
4243
#include "paimon/format/parquet/parquet_format_writer.h"
4344
#include "paimon/fs/file_system.h"
4445
#include "paimon/fs/local/local_file_system.h"
@@ -73,7 +74,8 @@ class ParquetStatsExtractorTest : public ::testing::Test {
7374
::parquet::WriterProperties::Builder builder;
7475
builder.enable_store_decimal_as_integer();
7576
ASSERT_OK_AND_ASSIGN(auto format_writer,
76-
ParquetFormatWriter::Create(out, arrow_schema, builder.build(), pool));
77+
ParquetFormatWriter::Create(out, arrow_schema, builder.build(), pool,
78+
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE));
7779
auto array = arrow::ipc::internal::json::ArrayFromJSON(struct_type, input).ValueOrDie();
7880
auto arrow_array = std::make_unique<ArrowArray>();
7981
ASSERT_TRUE(arrow::ExportArray(*array, arrow_array.get()).ok());
@@ -272,7 +274,8 @@ TEST_F(ParquetStatsExtractorTest, TestNullForAllType) {
272274
::parquet::WriterProperties::Builder builder;
273275
builder.enable_store_decimal_as_integer();
274276
ASSERT_OK_AND_ASSIGN(auto format_writer,
275-
ParquetFormatWriter::Create(out, schema, builder.build(), arrow_pool));
277+
ParquetFormatWriter::Create(out, schema, builder.build(), arrow_pool,
278+
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE));
276279
auto src_array = std::dynamic_pointer_cast<arrow::StructArray>(
277280
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields}), R"([
278281
[null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null,null]

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,11 @@ Result<std::unique_ptr<FormatWriter>> ParquetWriterBuilder::Build(
4242
const std::shared_ptr<OutputStream>& out, const std::string& compression) {
4343
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<::parquet::WriterProperties> writer_properties,
4444
PrepareWriterProperties(compression));
45-
return ParquetFormatWriter::Create(out, schema_, writer_properties, pool_);
45+
PAIMON_ASSIGN_OR_RAISE(uint64_t max_memory_use, OptionsUtils::GetValueFromMap<uint64_t>(
46+
options_, PARQUET_WRITER_MAX_MEMORY_USE,
47+
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE));
48+
49+
return ParquetFormatWriter::Create(out, schema_, writer_properties, pool_, max_memory_use);
4650
}
4751

4852
Result<std::shared_ptr<::parquet::WriterProperties>> ParquetWriterBuilder::PrepareWriterProperties(

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,13 +24,11 @@
2424

2525
#include "arrow/memory_pool.h"
2626
#include "arrow/type.h"
27-
#include "arrow/util/compression.h"
2827
#include "arrow/util/type_fwd.h"
2928
#include "paimon/common/utils/arrow/mem_utils.h"
3029
#include "paimon/format/format_writer.h"
3130
#include "paimon/format/writer_builder.h"
3231
#include "paimon/fs/file_system.h"
33-
#include "paimon/fs/file_system_factory.h"
3432
#include "paimon/memory/memory_pool.h"
3533
#include "paimon/result.h"
3634
#include "paimon/type_fwd.h"

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,8 @@ class PredicatePushdownTest : public ::testing::Test {
9393
auto writer_properties = builder.build();
9494
ASSERT_OK_AND_ASSIGN(
9595
auto format_writer,
96-
ParquetFormatWriter::Create(out, data_schema, writer_properties, arrow_pool_));
96+
ParquetFormatWriter::Create(out, data_schema, writer_properties, arrow_pool_,
97+
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE));
9798
ASSERT_OK(format_writer->AddBatch(data_arrow_array.get()));
9899
ASSERT_OK(format_writer->Finish());
99100
ASSERT_OK(out->Close());

0 commit comments

Comments
 (0)