From 1c4130759b5194d21205da0e00d3ddad2156b317 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Thu, 30 Jul 2026 21:06:38 +0800 Subject: [PATCH 1/2] feat(parquet): add size limit (in bytes) for row groups --- cpp/src/parquet/arrow/writer.cc | 20 +++++++++++++++++--- cpp/src/parquet/properties.h | 30 ++++++++++++++++++++++++++---- 2 files changed, 43 insertions(+), 7 deletions(-) diff --git a/cpp/src/parquet/arrow/writer.cc b/cpp/src/parquet/arrow/writer.cc index e0fbe308219c..81d575208e55 100644 --- a/cpp/src/parquet/arrow/writer.cc +++ b/cpp/src/parquet/arrow/writer.cc @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -456,10 +457,24 @@ class FileWriterImpl : public FileWriter { // Max number of rows allowed in a row group. const int64_t max_row_group_length = this->properties().max_row_group_length(); + // Max compressed byte size allowed in a row group. + const int64_t max_row_group_size = this->properties().max_row_group_size(); + const bool row_group_size_limited = + max_row_group_size != std::numeric_limits::max(); + + // Whether the current row group reached the row count or byte size limit. + auto row_group_full = [&]() { + return row_group_writer_->num_rows() >= max_row_group_length || + (row_group_size_limited && + row_group_writer_->total_compressed_bytes() + + row_group_writer_->total_compressed_bytes_written() + + row_group_writer_->estimated_buffered_stats().dict_bytes + >= max_row_group_size); + }; // Initialize a new buffered row group writer if necessary. if (row_group_writer_ == nullptr || !row_group_writer_->buffered() || - row_group_writer_->num_rows() >= max_row_group_length) { + row_group_full()) { RETURN_NOT_OK(NewBufferedRowGroup()); } @@ -501,8 +516,7 @@ class FileWriterImpl : public FileWriter { offset += batch_size; // Flush current row group writer and create a new writer if it is full. - if (row_group_writer_->num_rows() >= max_row_group_length && - offset < batch.num_rows()) { + if (row_group_full() && offset < batch.num_rows()) { RETURN_NOT_OK(NewBufferedRowGroup()); } } diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h index e2244a1176e3..cade6bea30f8 100644 --- a/cpp/src/parquet/properties.h +++ b/cpp/src/parquet/properties.h @@ -17,6 +17,7 @@ #pragma once +#include #include #include #include @@ -162,6 +163,8 @@ static constexpr bool DEFAULT_IS_DICTIONARY_ENABLED = true; static constexpr int64_t DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT = kDefaultDataPageSize; static constexpr int64_t DEFAULT_WRITE_BATCH_SIZE = 1024; static constexpr int64_t DEFAULT_MAX_ROW_GROUP_LENGTH = 1024 * 1024; +static constexpr int64_t DEFAULT_MAX_ROW_GROUP_SIZE = + std::numeric_limits::max(); static constexpr bool DEFAULT_ARE_STATISTICS_ENABLED = true; static constexpr int64_t DEFAULT_MAX_STATISTICS_SIZE = 4096; static constexpr Encoding::type DEFAULT_ENCODING = Encoding::UNKNOWN; @@ -367,6 +370,7 @@ class PARQUET_EXPORT WriterProperties { dictionary_pagesize_limit_(DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT), write_batch_size_(DEFAULT_WRITE_BATCH_SIZE), max_row_group_length_(DEFAULT_MAX_ROW_GROUP_LENGTH), + max_row_group_size_(DEFAULT_MAX_ROW_GROUP_SIZE), pagesize_(kDefaultDataPageSize), max_rows_per_page_(kDefaultMaxRowsPerPage), version_(ParquetVersion::PARQUET_2_6), @@ -383,6 +387,7 @@ class PARQUET_EXPORT WriterProperties { dictionary_pagesize_limit_(properties.dictionary_pagesize_limit()), write_batch_size_(properties.write_batch_size()), max_row_group_length_(properties.max_row_group_length()), + max_row_group_size_(properties.max_row_group_size()), pagesize_(properties.data_pagesize()), max_rows_per_page_(properties.max_rows_per_page()), version_(properties.version()), @@ -492,6 +497,18 @@ class PARQUET_EXPORT WriterProperties { return this; } + /// Specify the max row group size in compressed bytes. + /// Default unlimited. + /// + /// The limit is checked against the compressed pages accumulated in the + /// current row group, so the actual row group size may slightly exceed it. + /// Only effective for buffered row groups ( + /// parquet::arrow::FileWriter::WriteRecordBatch). + Builder* max_row_group_size(int64_t max_row_group_size) { + max_row_group_size_ = max_row_group_size; + return this; + } + /// Specify the data page size. /// Default 1MB. Builder* data_pagesize(int64_t pg_size) { @@ -899,8 +916,8 @@ class PARQUET_EXPORT WriterProperties { return std::shared_ptr(new WriterProperties( pool_, dictionary_pagesize_limit_, write_batch_size_, max_row_group_length_, - pagesize_, max_rows_per_page_, version_, created_by_, page_checksum_enabled_, - size_statistics_level_, std::move(file_encryption_properties_), + max_row_group_size_, pagesize_, max_rows_per_page_, version_, created_by_, + page_checksum_enabled_,size_statistics_level_, std::move(file_encryption_properties_), default_column_properties_, column_properties, data_page_version_, store_decimal_as_integer_, std::move(sorting_columns_), content_defined_chunking_enabled_, content_defined_chunking_options_)); @@ -913,6 +930,7 @@ class PARQUET_EXPORT WriterProperties { int64_t dictionary_pagesize_limit_; int64_t write_batch_size_; int64_t max_row_group_length_; + int64_t max_row_group_size_; int64_t pagesize_; int64_t max_rows_per_page_; ParquetVersion::type version_; @@ -949,6 +967,8 @@ class PARQUET_EXPORT WriterProperties { inline int64_t max_row_group_length() const { return max_row_group_length_; } + inline int64_t max_row_group_size() const { return max_row_group_size_; } + inline int64_t data_pagesize() const { return pagesize_; } inline int64_t max_rows_per_page() const { return max_rows_per_page_; } @@ -1078,8 +1098,8 @@ class PARQUET_EXPORT WriterProperties { private: explicit WriterProperties( MemoryPool* pool, int64_t dictionary_pagesize_limit, int64_t write_batch_size, - int64_t max_row_group_length, int64_t pagesize, int64_t max_rows_per_page, - ParquetVersion::type version, const std::string& created_by, + int64_t max_row_group_length, int64_t max_row_group_size, int64_t pagesize, + int64_t max_rows_per_page, ParquetVersion::type version, const std::string& created_by, bool page_write_checksum_enabled, SizeStatisticsLevel size_statistics_level, std::shared_ptr file_encryption_properties, const ColumnProperties& default_column_properties, @@ -1091,6 +1111,7 @@ class PARQUET_EXPORT WriterProperties { dictionary_pagesize_limit_(dictionary_pagesize_limit), write_batch_size_(write_batch_size), max_row_group_length_(max_row_group_length), + max_row_group_size_(max_row_group_size), pagesize_(pagesize), max_rows_per_page_(max_rows_per_page), parquet_data_page_version_(data_page_version), @@ -1110,6 +1131,7 @@ class PARQUET_EXPORT WriterProperties { int64_t dictionary_pagesize_limit_; int64_t write_batch_size_; int64_t max_row_group_length_; + int64_t max_row_group_size_; int64_t pagesize_; int64_t max_rows_per_page_; ParquetDataPageVersion parquet_data_page_version_; From 4acdeeee02e73cdaa43c47a1f5a5ffdf4a1b2138 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Thu, 30 Jul 2026 21:38:22 +0800 Subject: [PATCH 2/2] add test cases --- .../parquet/arrow/arrow_reader_writer_test.cc | 50 +++++++++++++++++++ cpp/src/parquet/arrow/writer.cc | 10 ++-- cpp/src/parquet/properties.h | 23 +++++---- cpp/src/parquet/properties_test.cc | 2 + 4 files changed, 69 insertions(+), 16 deletions(-) diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index 8735aea731ce..66bed708f3a0 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -5888,6 +5888,56 @@ TEST(TestArrowReadWrite, WriteRecordBatchNotProduceEmptyRowGroup) { } } +TEST(TestArrowReadWrite, WriteRecordBatchRespectsMaxRowGroupSize) { + // Row groups should be rolled over once the compressed bytes accumulated + // in the current row group reach WriterProperties::max_row_group_size(). + auto pool = ::arrow::default_memory_pool(); + auto sink = CreateOutputStream(); + // Use a small byte size limit with the default (large) row count limit so + // that only the byte size limit takes effect. Use a small data page size + // so that buffered values are flushed to pages (and thus counted by the + // size check) between batches. + auto writer_properties = WriterProperties::Builder() + .max_row_group_size(4 * 1024) + ->disable_dictionary() + ->data_pagesize(1024) + ->build(); + auto arrow_writer_properties = default_arrow_writer_properties(); + + // Prepare schema + auto schema = ::arrow::schema({::arrow::field("a", ::arrow::int64())}); + std::shared_ptr parquet_schema; + ASSERT_OK_NO_THROW(ToParquetSchema(schema.get(), *writer_properties, + *arrow_writer_properties, &parquet_schema)); + auto schema_node = std::static_pointer_cast(parquet_schema->schema_root()); + + auto gen = ::arrow::random::RandomArrayGenerator(/*seed=*/42); + + // Create writer to write data via RecordBatch. + auto writer = ParquetFileWriter::Open(sink, schema_node, writer_properties); + std::unique_ptr arrow_writer; + ASSERT_OK(FileWriter::Make(pool, std::move(writer), schema, arrow_writer_properties, + &arrow_writer)); + // Each batch holds 1000 int64 values (~8KB uncompressed), exceeding the + // 4KB limit on its own, so every subsequent WriteRecordBatch call should + // start a new row group. + constexpr int kNumBatches = 4; + constexpr int64_t kBatchRows = 1000; + for (int i = 0; i < kNumBatches; ++i) { + auto record_batch = gen.BatchOf({::arrow::field("a", ::arrow::int64())}, + /*length=*/kBatchRows); + ASSERT_OK_NO_THROW(arrow_writer->WriteRecordBatch(*record_batch)); + } + ASSERT_OK_NO_THROW(arrow_writer->Close()); + ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish()); + + auto file_metadata = arrow_writer->metadata(); + ASSERT_EQ(kNumBatches, file_metadata->num_row_groups()); + for (int i = 0; i < file_metadata->num_row_groups(); ++i) { + EXPECT_EQ(kBatchRows, file_metadata->RowGroup(i)->num_rows()); + } +} + TEST(TestArrowReadWrite, MultithreadedWrite) { const int num_columns = 20; const int num_rows = 1000; diff --git a/cpp/src/parquet/arrow/writer.cc b/cpp/src/parquet/arrow/writer.cc index 81d575208e55..185f047f88d7 100644 --- a/cpp/src/parquet/arrow/writer.cc +++ b/cpp/src/parquet/arrow/writer.cc @@ -465,11 +465,11 @@ class FileWriterImpl : public FileWriter { // Whether the current row group reached the row count or byte size limit. auto row_group_full = [&]() { return row_group_writer_->num_rows() >= max_row_group_length || - (row_group_size_limited && - row_group_writer_->total_compressed_bytes() + - row_group_writer_->total_compressed_bytes_written() + - row_group_writer_->estimated_buffered_stats().dict_bytes - >= max_row_group_size); + (row_group_size_limited && + row_group_writer_->total_compressed_bytes() + + row_group_writer_->total_compressed_bytes_written() + + row_group_writer_->estimated_buffered_stats().dict_bytes >= + max_row_group_size); }; // Initialize a new buffered row group writer if necessary. diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h index cade6bea30f8..0409c4152931 100644 --- a/cpp/src/parquet/properties.h +++ b/cpp/src/parquet/properties.h @@ -163,8 +163,7 @@ static constexpr bool DEFAULT_IS_DICTIONARY_ENABLED = true; static constexpr int64_t DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT = kDefaultDataPageSize; static constexpr int64_t DEFAULT_WRITE_BATCH_SIZE = 1024; static constexpr int64_t DEFAULT_MAX_ROW_GROUP_LENGTH = 1024 * 1024; -static constexpr int64_t DEFAULT_MAX_ROW_GROUP_SIZE = - std::numeric_limits::max(); +static constexpr int64_t DEFAULT_MAX_ROW_GROUP_SIZE = std::numeric_limits::max(); static constexpr bool DEFAULT_ARE_STATISTICS_ENABLED = true; static constexpr int64_t DEFAULT_MAX_STATISTICS_SIZE = 4096; static constexpr Encoding::type DEFAULT_ENCODING = Encoding::UNKNOWN; @@ -499,7 +498,7 @@ class PARQUET_EXPORT WriterProperties { /// Specify the max row group size in compressed bytes. /// Default unlimited. - /// + /// /// The limit is checked against the compressed pages accumulated in the /// current row group, so the actual row group size may slightly exceed it. /// Only effective for buffered row groups ( @@ -916,11 +915,12 @@ class PARQUET_EXPORT WriterProperties { return std::shared_ptr(new WriterProperties( pool_, dictionary_pagesize_limit_, write_batch_size_, max_row_group_length_, - max_row_group_size_, pagesize_, max_rows_per_page_, version_, created_by_, - page_checksum_enabled_,size_statistics_level_, std::move(file_encryption_properties_), - default_column_properties_, column_properties, data_page_version_, - store_decimal_as_integer_, std::move(sorting_columns_), - content_defined_chunking_enabled_, content_defined_chunking_options_)); + max_row_group_size_, pagesize_, max_rows_per_page_, version_, created_by_, + page_checksum_enabled_, size_statistics_level_, + std::move(file_encryption_properties_), default_column_properties_, + column_properties, data_page_version_, store_decimal_as_integer_, + std::move(sorting_columns_), content_defined_chunking_enabled_, + content_defined_chunking_options_)); } private: @@ -1098,9 +1098,10 @@ class PARQUET_EXPORT WriterProperties { private: explicit WriterProperties( MemoryPool* pool, int64_t dictionary_pagesize_limit, int64_t write_batch_size, - int64_t max_row_group_length, int64_t max_row_group_size, int64_t pagesize, - int64_t max_rows_per_page, ParquetVersion::type version, const std::string& created_by, - bool page_write_checksum_enabled, SizeStatisticsLevel size_statistics_level, + int64_t max_row_group_length, int64_t max_row_group_size, int64_t pagesize, + int64_t max_rows_per_page, ParquetVersion::type version, + const std::string& created_by, bool page_write_checksum_enabled, + SizeStatisticsLevel size_statistics_level, std::shared_ptr file_encryption_properties, const ColumnProperties& default_column_properties, const std::unordered_map& column_properties, diff --git a/cpp/src/parquet/properties_test.cc b/cpp/src/parquet/properties_test.cc index 324ea2026a9e..e91a72f609da 100644 --- a/cpp/src/parquet/properties_test.cc +++ b/cpp/src/parquet/properties_test.cc @@ -288,6 +288,7 @@ TEST_P(WriterPropertiesTest, RoundTripThroughBuilder) { properties->file_encryption_properties()); ASSERT_EQ(round_tripped->max_rows_per_page(), properties->max_rows_per_page()); ASSERT_EQ(round_tripped->max_row_group_length(), properties->max_row_group_length()); + ASSERT_EQ(round_tripped->max_row_group_size(), properties->max_row_group_size()); ASSERT_EQ(round_tripped->memory_pool(), properties->memory_pool()); ASSERT_EQ(round_tripped->page_checksum_enabled(), properties->page_checksum_enabled()); ASSERT_EQ(round_tripped->size_statistics_level(), properties->size_statistics_level()); @@ -352,6 +353,7 @@ std::vector writer_properties_test_cases() { builder.dictionary_pagesize_limit(DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT - 1); builder.write_batch_size(DEFAULT_WRITE_BATCH_SIZE - 1); builder.max_row_group_length(DEFAULT_MAX_ROW_GROUP_LENGTH - 1); + builder.max_row_group_size(DEFAULT_MAX_ROW_GROUP_SIZE - 1); builder.data_pagesize(kDefaultDataPageSize - 1); builder.max_rows_per_page(kDefaultMaxRowsPerPage - 1); builder.data_page_version(ParquetDataPageVersion::V2);