feat: use buffered path for WriteTable if max row group size is set

Author: zhf999

cpp/src/parquet/arrow/arrow_reader_writer_test.cc

@@ -5894,13 +5894,12 @@ TEST(TestArrowReadWrite, WriteRecordBatchRespectsMaxRowGroupSize) {
   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.
+  // 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()
-                               ->data_pagesize(1024)
                                ->build();
   auto arrow_writer_properties = default_arrow_writer_properties();
 
@@ -5938,6 +5937,87 @@ TEST(TestArrowReadWrite, WriteRecordBatchRespectsMaxRowGroupSize) {
   }
 }
 
+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;
+  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();
+    // The size estimate is conservative, so row groups stay under the limit.
+    EXPECT_LE(row_group_metadata->total_compressed_size(), kMaxRowGroupSize);
+  }
+  // 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

@@ -432,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))
@@ -450,26 +467,38 @@ 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 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();
 
+    // Estimated size of the data accumulated in the current row group.
+    auto estimated_row_group_size = [&]() {
+      const auto buffered = row_group_writer_->estimated_buffered_stats();
+      return 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;
+    };
+
     // 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);
+      return row_group_writer_->num_rows() >= max_rows_per_row_group ||
+             (row_group_size_limited && estimated_row_group_size() >= max_row_group_size);
     };
 
     // Initialize a new buffered row group writer if necessary.
@@ -510,7 +539,7 @@ 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;

cpp/src/parquet/properties.h

@@ -496,13 +496,19 @@ class PARQUET_EXPORT WriterProperties {
       return this;
     }
 
-    /// Specify the max row group size in compressed bytes.
+    /// Specify the max row group size in 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).
+    /// The limit is checked between writes against an estimate of the data
+    /// accumulated in the current row group, which combines the size of the
+    /// compressed pages with an uncompressed estimate of the values still
+    /// buffered by the column encoders. The estimate is therefore approximate
+    /// and errs on the conservative side.
+    /// Row groups are buffered in memory while the limit is enforced, so
+    /// enabling this option increases the memory footprint of the writer.
+    /// It is honoured by parquet::arrow::FileWriter::WriteRecordBatch and
+    /// parquet::arrow::FileWriter::WriteTable; row groups created explicitly
+    /// 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;