Skip to content

Commit eacebd4

Browse files
zjw1111zjw1111
committed
fix
1 parent a45c1bf commit eacebd4

4 files changed

Lines changed: 65 additions & 26 deletions

File tree

‎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
@@ -111,8 +111,8 @@ class ParquetFileBatchReaderTest : public ::testing::Test,
111111
enable_dictionary ? builder.enable_dictionary() : builder.disable_dictionary();
112112
auto writer_properties = builder.build();
113113
ASSERT_OK_AND_ASSIGN(auto format_writer, ParquetFormatWriter::Create(
114-
out, arrow_schema, writer_properties, pool_,
115-
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE));
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: 3 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";
@@ -38,8 +38,8 @@ static inline const char PARQUET_WRITER_MAX_MEMORY_USE[] = "parquet.writer.max.m
3838
static constexpr uint64_t DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE = 512 * 1024 * 1024; // 512MB
3939

4040
// read
41-
static inline const char PARQUET_READ_USE_THREADS[] = "parquet.read.use-threads";
42-
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;
4343
static inline const char PARQUET_READ_CACHE_OPTION_LAZY[] = "parquet.read.cache-option.lazy";
4444
static inline const char PARQUET_READ_CACHE_OPTION_PREFETCH_LIMIT[] =
4545
"parquet.read.cache-option.prefetch-limit";

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

Lines changed: 56 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -130,15 +130,16 @@ class ParquetFormatWriterTest : public ::testing::Test {
130130
ASSERT_OK(format_writer->AddBatch(batch->GetData()));
131131
}
132132

133-
void CheckResult(const std::string& file_path, int32_t row_count) const {
133+
void CheckResult(const std::string& file_path, int32_t row_count,
134+
int32_t row_group_count) const {
134135
auto file = arrow::io::ReadableFile::Open(file_path, arrow_pool_.get());
135136
ASSERT_TRUE(file.ok());
136137
std::unique_ptr<::parquet::arrow::FileReader> reader;
137138
auto status = ::parquet::arrow::OpenFile(file.ValueOrDie(), arrow_pool_.get(), &reader);
138139
ASSERT_TRUE(status.ok()) << status.ToString();
139140
const ::parquet::FileMetaData* metadata = reader->parquet_reader()->metadata().get();
140141
const ::parquet::SchemaDescriptor* schema = metadata->schema();
141-
ASSERT_EQ(metadata->num_row_groups(), 1);
142+
ASSERT_EQ(metadata->num_row_groups(), row_group_count);
142143
ASSERT_EQ(schema->num_columns(), 3);
143144
ASSERT_EQ(metadata->num_rows(), row_count);
144145
ASSERT_EQ("col1", schema->Column(0)->name());
@@ -205,8 +206,8 @@ TEST_F(ParquetFormatWriterTest, TestWriteWithVariousBatchSize) {
205206
auto writer_properties = builder.build();
206207
ASSERT_OK_AND_ASSIGN(
207208
auto format_writer,
208-
ParquetFormatWriter::Create(out, arrow_schema, writer_properties, arrow_pool_,
209-
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE));
209+
ParquetFormatWriter::Create(out, arrow_schema, writer_properties,
210+
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, arrow_pool_));
210211
auto array = PrepareArray(struct_type, record_batch_size);
211212
auto arrow_array = std::make_unique<ArrowArray>();
212213
ASSERT_TRUE(arrow::ExportArray(*array, arrow_array.get()).ok());
@@ -219,7 +220,7 @@ TEST_F(ParquetFormatWriterTest, TestWriteWithVariousBatchSize) {
219220
ASSERT_OK(format_writer->Finish());
220221
ASSERT_OK(out->Flush());
221222
ASSERT_OK(out->Close());
222-
CheckResult(file_path, record_batch_size);
223+
CheckResult(file_path, record_batch_size, /*row_group_count=*/1);
223224
}
224225
}
225226
}
@@ -239,8 +240,8 @@ TEST_F(ParquetFormatWriterTest, TestWriteMultipleTimes) {
239240
auto writer_properties = builder.build();
240241
ASSERT_OK_AND_ASSIGN(
241242
std::shared_ptr<ParquetFormatWriter> format_writer,
242-
ParquetFormatWriter::Create(out, arrow_schema, writer_properties, arrow_pool_,
243-
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE));
243+
ParquetFormatWriter::Create(out, arrow_schema, writer_properties,
244+
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, arrow_pool_));
244245

245246
// add batch first time, 6 rows
246247
AddRecordBatchOnce(format_writer, struct_type, 6, 0);
@@ -264,7 +265,7 @@ TEST_F(ParquetFormatWriterTest, TestWriteMultipleTimes) {
264265
ASSERT_OK(format_writer->Finish());
265266
ASSERT_OK(out->Flush());
266267
ASSERT_OK(out->Close());
267-
CheckResult(file_path, /*row_count=*/37);
268+
CheckResult(file_path, /*row_count=*/37, /*row_group_count=*/1);
268269
auto metrics = format_writer->GetWriterMetrics();
269270
ASSERT_OK_AND_ASSIGN(uint64_t counter, metrics->GetCounter(ParquetMetrics::WRITE_RECORD_COUNT));
270271
ASSERT_EQ(37, counter);
@@ -282,8 +283,8 @@ TEST_F(ParquetFormatWriterTest, TestGetEstimateLength) {
282283
auto writer_properties = builder.build();
283284
ASSERT_OK_AND_ASSIGN(
284285
std::shared_ptr<ParquetFormatWriter> format_writer,
285-
ParquetFormatWriter::Create(out, arrow_schema, writer_properties, arrow_pool_,
286-
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE));
286+
ParquetFormatWriter::Create(out, arrow_schema, writer_properties,
287+
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, arrow_pool_));
287288

288289
// add batch first time, 1 row
289290
AddRecordBatchOnce(format_writer, struct_type, 1, 0);
@@ -302,7 +303,7 @@ TEST_F(ParquetFormatWriterTest, TestGetEstimateLength) {
302303
}
303304

304305
TEST_F(ParquetFormatWriterTest, TestMemoryControl) {
305-
auto check_result = [&](bool all_null_value, uint64_t max_memory_use) {
306+
auto run = [&](bool all_null_value, uint64_t max_memory_use) {
306307
ASSERT_OK_AND_ASSIGN(
307308
std::unique_ptr<FileFormat> file_format,
308309
FileFormatFactory::Get(
@@ -348,10 +349,49 @@ TEST_F(ParquetFormatWriterTest, TestMemoryControl) {
348349
ASSERT_GT(actual_max_mem, max_memory_use);
349350
ASSERT_LT(actual_max_mem, max_memory_use * 1.5); // allow 50% overhead
350351
};
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
352+
run(/*all_null_value=*/true, /*max_memory_use=*/20 * 1024 * 1024); // 20MB
353+
run(/*all_null_value=*/true, /*max_memory_use=*/40 * 1024 * 1024); // 40MB
354+
run(/*all_null_value=*/false, /*max_memory_use=*/20 * 1024 * 1024); // 20MB
355+
run(/*all_null_value=*/false, /*max_memory_use=*/40 * 1024 * 1024); // 40MB
356+
}
357+
358+
TEST_F(ParquetFormatWriterTest, TestMemoryControlForCheckRowGroupCount) {
359+
auto run = [&](int32_t write_times) {
360+
ASSERT_OK_AND_ASSIGN(
361+
std::unique_ptr<FileFormat> file_format,
362+
FileFormatFactory::Get("parquet", {{Options::FILE_FORMAT, "parquet"},
363+
{Options::MANIFEST_FORMAT, "parquet"},
364+
{"parquet.writer.max.memory.use", "1"}}));
365+
366+
auto schema_pair = PrepareArrowSchema();
367+
const auto& arrow_schema = schema_pair.first;
368+
const auto& struct_type = schema_pair.second;
369+
int32_t batch_size = 4096;
370+
std::string file_path =
371+
PathUtil::JoinPath(dir_->Str(), std::to_string(write_times) + ".parquet");
372+
373+
auto c_schema = std::make_unique<::ArrowSchema>();
374+
ASSERT_TRUE(arrow::ExportSchema(*arrow_schema, c_schema.get()).ok());
375+
ASSERT_OK_AND_ASSIGN(auto writer_builder,
376+
file_format->CreateWriterBuilder(c_schema.get(), batch_size));
377+
ASSERT_OK_AND_ASSIGN(std::shared_ptr<OutputStream> out,
378+
fs_->Create(file_path, /*overwrite=*/false));
379+
ASSERT_OK_AND_ASSIGN(std::shared_ptr<FormatWriter> writer,
380+
writer_builder->Build(out, "uncompressed"));
381+
382+
for (int32_t i = 0; i < write_times; ++i) {
383+
AddRecordBatchOnce(writer, struct_type, 10, i * 10);
384+
}
385+
386+
ASSERT_OK(writer->Flush());
387+
ASSERT_OK(writer->Finish());
388+
ASSERT_OK(out->Flush());
389+
ASSERT_OK(out->Close());
390+
CheckResult(file_path, /*row_count=*/write_times * 10, /*row_group_count=*/write_times);
391+
};
392+
run(/*write_times=*/1);
393+
run(/*write_times=*/2);
394+
run(/*write_times=*/5);
355395
}
356396

357397
TEST_F(ParquetFormatWriterTest, TestTimestampType) {
@@ -363,7 +403,6 @@ TEST_F(ParquetFormatWriterTest, TestTimestampType) {
363403
arrow::field("ts_nano", arrow::timestamp(arrow::TimeUnit::NANO)),
364404
arrow::field("ts_utc1", arrow::timestamp(arrow::TimeUnit::SECOND, timezone)),
365405
arrow::field("ts_utc2", arrow::timestamp(arrow::TimeUnit::MICRO, timezone))};
366-
[[maybe_unused]] auto test = new arrow::Schema(fields); // for test
367406

368407
std::string file_path = PathUtil::JoinPath(dir_->Str(), "timezone.parquet");
369408
ASSERT_OK_AND_ASSIGN(std::shared_ptr<OutputStream> out,
@@ -373,7 +412,7 @@ TEST_F(ParquetFormatWriterTest, TestTimestampType) {
373412
ASSERT_OK_AND_ASSIGN(
374413
std::shared_ptr<ParquetFormatWriter> format_writer,
375414
ParquetFormatWriter::Create(out, std::make_shared<arrow::Schema>(fields), writer_properties,
376-
arrow_pool_, DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE));
415+
DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, arrow_pool_));
377416

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

0 commit comments

Comments
 (0)