Skip to content

GH-48476: [C++] [Parquet] Add max_row_group_size writer property to limit row groups by compressed byte size #50743

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

Open
zhf999 wants to merge 2 commits into
base: main
Choose a base branch
from
+100 −11
Open

GH-48476: [C++] [Parquet] Add max_row_group_size writer property to limit row groups by compressed byte size #50743

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1c41307
feat(parquet): add size limit (in bytes) for row groups
zhf999 Jul 30, 2026
4acdeee
add test cases
zhf999 Jul 30, 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
50 changes: 50 additions & 0 deletions cpp/src/parquet/arrow/arrow_reader_writer_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<SchemaDescriptor> parquet_schema;
ASSERT_OK_NO_THROW(ToParquetSchema(schema.get(), *writer_properties,
*arrow_writer_properties, &parquet_schema));
auto schema_node = std::static_pointer_cast<GroupNode>(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<FileWriter> 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;
Expand Down
20 changes: 17 additions & 3 deletions cpp/src/parquet/arrow/writer.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

#include <algorithm>
#include <deque>
#include <limits>
#include <memory>
#include <string>
#include <type_traits>
Expand Down Expand Up @@ -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<int64_t>::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());
}

Expand Down Expand Up @@ -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());
}
}
Expand Down
39 changes: 31 additions & 8 deletions cpp/src/parquet/properties.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

#pragma once

#include <limits>
#include <memory>
#include <optional>
#include <string>
Expand Down Expand Up @@ -162,6 +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<int64_t>::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;
Expand Down Expand Up @@ -367,6 +369,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),
Expand All @@ -383,6 +386,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()),
Expand Down Expand Up @@ -492,6 +496,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).

@wgtmac wgtmac Jul 30, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why only for WriteRecordBatch?

@zhf999 zhf999 Jul 31, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WriteRecordBatch writes into a buffered row group, which accumulates across calls. That makes the byte size observable between writes, so we can simply stop adding to the current row group once the accumulated compressed size reaches the limit.

While WriteTable uses the non-buffered path, where each call maps to exactly one row group whose row count is fixed up front by the user-supplied chunk_size. Enforcing a byte limit there would require predicting the compressed size before writing, e.g. by estimating an average row size from the previous row group. That's inherently approximate, especially since chunk_size is an explicit contract from the caller.

I wonder if estimation is acceptable in WriteTable?

@wgtmac wgtmac Jul 31, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IIRC, WriteTable also splits the table into several row groups if max_row_group_length is reached.

@zhf999 zhf999 Jul 31, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

max_row_group_length can be applied there because row counts are known before writing — it just clamps chunk_size. The byte size simply isn't knowable at that point, which is the asymmetry.

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) {
Expand Down Expand Up @@ -899,11 +915,12 @@ class PARQUET_EXPORT WriterProperties {

return std::shared_ptr<WriterProperties>(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_),
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:
Expand All @@ -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_;
Expand Down Expand Up @@ -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_; }
Expand Down Expand Up @@ -1078,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 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<FileEncryptionProperties> file_encryption_properties,
const ColumnProperties& default_column_properties,
const std::unordered_map<std::string, ColumnProperties>& column_properties,
Expand All @@ -1091,6 +1112,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),
Expand All @@ -1110,6 +1132,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_;
Expand Down
2 changes: 2 additions & 0 deletions cpp/src/parquet/properties_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand Down Expand Up @@ -352,6 +353,7 @@ std::vector<WriterPropertiesTestCase> 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);
Expand Down
Loading