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

cpp/src/parquet/arrow/arrow_reader_writer_test.cc

@@ -5888,6 +5888,138 @@ 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. The default data page size is
+  // kept, so this also covers values that are still buffered by the column
+  // encoders rather than flushed into pages.
+  auto writer_properties = WriterProperties::Builder()
+                               .max_row_group_size(4 * 1024)
+                               ->disable_dictionary()
+                               ->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, WriteTableRespectsMaxRowGroupSize) {
+  // When a byte size limit is set, WriteTable switches to buffered row groups
+  // and feeds the table in small batches, so a chunk_size covering the whole
+  // table is still split into row groups bounded by the limit.
+  constexpr int64_t kMaxRowGroupSize = 16 * 1024;
+  // The limit is enforced against an approximate size estimate, so row groups
+  // are only expected to stay in the vicinity of it.
+  constexpr int64_t kSlack = 8 * 1024;
+  constexpr int kNumRows = 8000;
+
+  std::shared_ptr<Table> table;
+  ASSERT_NO_FATAL_FAILURE(
+      MakeDoubleTable(/*num_columns=*/1, kNumRows, /*nchunks=*/1, &table));
+
+  auto sink = CreateOutputStream();
+  // The default data page size (1MB) is much larger than the limit here, so
+  // this only works if values still buffered by the column encoders count
+  // towards the row group size.
+  auto writer_properties = WriterProperties::Builder()
+                               .max_row_group_size(kMaxRowGroupSize)
+                               ->write_batch_size(256)
+                               ->disable_dictionary()
+                               ->build();
+  ASSERT_OK_NO_THROW(WriteTable(*table, ::arrow::default_memory_pool(), sink,
+                                /*chunk_size=*/kNumRows, writer_properties));
+  ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish());
+
+  auto reader = ParquetFileReader::Open(std::make_shared<BufferReader>(buffer));
+  auto file_metadata = reader->metadata();
+  // The limit must have split the single chunk into multiple row groups.
+  ASSERT_GT(file_metadata->num_row_groups(), 1);
+
+  int64_t total_rows = 0;
+  for (int i = 0; i < file_metadata->num_row_groups(); ++i) {
+    auto row_group_metadata = file_metadata->RowGroup(i);
+    total_rows += row_group_metadata->num_rows();
+    EXPECT_LE(row_group_metadata->total_compressed_size(), kMaxRowGroupSize + kSlack);
+  }
+  // All rows are written exactly once.
+  EXPECT_EQ(kNumRows, total_rows);
+}
+
+TEST(TestArrowReadWrite, WriteTableUnlimitedRowGroupSize) {
+  // Without an explicit byte size limit, chunk_size alone decides the row group
+  // boundaries and row groups are not buffered.
+  constexpr int kNumRows = 2000;
+  std::shared_ptr<Table> table;
+  ASSERT_NO_FATAL_FAILURE(
+      MakeDoubleTable(/*num_columns=*/1, kNumRows, /*nchunks=*/1, &table));
+
+  auto sink = CreateOutputStream();
+  ASSERT_OK_NO_THROW(WriteTable(*table, ::arrow::default_memory_pool(), sink,
+                                /*chunk_size=*/1000, default_writer_properties()));
+  ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish());
+
+  auto reader = ParquetFileReader::Open(std::make_shared<BufferReader>(buffer));
+  ASSERT_EQ(2, reader->metadata()->num_row_groups());
+  EXPECT_EQ(1000, reader->metadata()->RowGroup(0)->num_rows());
+  EXPECT_EQ(1000, reader->metadata()->RowGroup(1)->num_rows());
+}
+
+TEST(TestArrowReadWrite, WriteTableMaxRowGroupSizeRoundTrip) {
+  // The data must survive the switch to the buffered write path unchanged.
+  constexpr int kNumRows = 3000;
+  std::shared_ptr<Table> table;
+  ASSERT_NO_FATAL_FAILURE(
+      MakeDoubleTable(/*num_columns=*/3, kNumRows, /*nchunks=*/1, &table));
+
+  auto sink = CreateOutputStream();
+  auto writer_properties =
+      WriterProperties::Builder().max_row_group_size(16 * 1024)->build();
+  ASSERT_OK_NO_THROW(WriteTable(*table, ::arrow::default_memory_pool(), sink,
+                                /*chunk_size=*/kNumRows, writer_properties));
+  ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish());
+
+  ASSERT_OK_AND_ASSIGN(auto reader, OpenFile(std::make_shared<BufferReader>(buffer),
+                                             ::arrow::default_memory_pool()));
+  std::shared_ptr<Table> result;
+  ASSERT_OK_NO_THROW(reader->ReadTable(&result));
+  ASSERT_OK(result->ValidateFull());
+  AssertTablesEqual(*table, *result, /*same_chunk_layout=*/false);
+}
+
 TEST(TestArrowReadWrite, MultithreadedWrite) {
   const int num_columns = 20;
   const int num_rows = 1000;

cpp/src/parquet/arrow/writer.cc

@@ -19,6 +19,7 @@
 
 #include <algorithm>
 #include <deque>
+#include <limits>
 #include <memory>
 #include <string>
 #include <type_traits>
@@ -431,6 +432,23 @@ class FileWriterImpl : public FileWriter {
           WriteRowGroup(0, 0).OrElse([&](auto&&) { PARQUET_IGNORE_NOT_OK(Close()); }));
     }
 
+    // If max_row_group_size is set, use buffered path to write row groups.
+    if (this->properties().max_row_group_size() != std::numeric_limits<int64_t>::max()) {
+      ::arrow::TableBatchReader reader(table);
+      reader.set_chunksize(std::min(chunk_size, this->properties().write_batch_size()));
+      while (true) {
+        std::shared_ptr<RecordBatch> batch;
+        RETURN_NOT_OK(reader.ReadNext(&batch));
+        if (batch == nullptr) {
+          break;
+        }
+        RETURN_NOT_OK(WriteRecordBatchBuffered(*batch, chunk_size).OrElse([&](auto&&) {
+          PARQUET_IGNORE_NOT_OK(Close());
+        }));
+      }
+      return Status::OK();
+    }
+
     for (int chunk = 0; chunk * chunk_size < table.num_rows(); chunk++) {
       int64_t offset = chunk * chunk_size;
       RETURN_NOT_OK(WriteRowGroup(offset, std::min(chunk_size, table.num_rows() - offset))
@@ -449,17 +467,44 @@ class FileWriterImpl : public FileWriter {
   }
 
   Status WriteRecordBatch(const RecordBatch& batch) override {
+    // Checked up front because reading the properties of a closed file throws.
+    RETURN_NOT_OK(CheckClosed());
+    return WriteRecordBatchBuffered(batch, this->properties().max_row_group_length());
+  }
+
+  // Writes `batch` into buffered row groups, starting a new row group whenever
+  // the current one reaches `max_rows_per_row_group` rows or the max row group
+  // size in bytes configured on the writer properties.
+  Status WriteRecordBatchBuffered(const RecordBatch& batch,
+                                  int64_t max_rows_per_row_group) {
     RETURN_NOT_OK(CheckClosed());
     if (batch.num_rows() == 0) {
       return Status::OK();
     }
 
-    // Max number of rows allowed in a row group.
-    const int64_t max_row_group_length = this->properties().max_row_group_length();
+    // Max 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.
+    // Values not yet formed into a page are still held by the column encoders,
+    // so they are counted through EstimatedBufferedStats(); those estimates are
+    // uncompressed and exclude future page headers, hence only approximate.
+    auto row_group_full = [&]() {
+      const auto buffered = row_group_writer_->estimated_buffered_stats();
+      const int64_t estimated_size = row_group_writer_->total_compressed_bytes() +
+                                     row_group_writer_->total_compressed_bytes_written() +
+                                     buffered.value_bytes + buffered.def_level_bytes +
+                                     buffered.rep_level_bytes + buffered.dict_bytes;
+
+      return row_group_writer_->num_rows() >= max_rows_per_row_group ||
+             (row_group_size_limited && estimated_size >= 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());
     }
 
@@ -495,14 +540,13 @@ class FileWriterImpl : public FileWriter {
     int64_t offset = 0;
     while (offset < batch.num_rows()) {
       const int64_t batch_size =
-          std::min(max_row_group_length - row_group_writer_->num_rows(),
+          std::min(max_rows_per_row_group - row_group_writer_->num_rows(),
                    batch.num_rows() - offset);
       RETURN_NOT_OK(WriteBatch(offset, batch_size));
       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());
       }
     }

cpp/src/parquet/properties.h

@@ -17,6 +17,7 @@
 
 #pragma once
 
+#include <limits>
 #include <memory>
 #include <optional>
 #include <string>
@@ -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;
@@ -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),
@@ -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()),
@@ -492,6 +496,19 @@ class PARQUET_EXPORT WriterProperties {
       return this;
     }
 
+    /// Specify the max row group size in bytes.
+    /// Default unlimited.
+    ///
+    /// The limit is only approximate, as it is checked between writes against an
+    /// estimate of the accumulated data. Row groups are buffered in memory to
+    /// enforce it, which increases the memory footprint of the writer.
+    /// Row groups created through NewRowGroup()/WriteColumnChunk() are not
+    /// affected.
+    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,11 +916,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:
@@ -913,6 +931,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 +968,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,9 +1099,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,
@@ -1091,6 +1113,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 +1133,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_;

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<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);