Repository: apache/arrow
File: cpp/src/parquet/arrow/arrow_reader_writer_test.cc (C++)

Code before:
      [6,    {"b1": null, "b2": "-22" }, "alfa"],
      [7,    {"b1": 0,    "b2": "-1"  }, "beta"],
      [8,    {"b1": null, "b2": null  }, "beta"],
      [9,    {"b1": 1,    "b2": "0"   }, null  ],
      [null, {"b1": null, "b2": ""    }, "gama"],
      [11,   {"b1": 2,    "b2": "1234"}, "foo" ],
      [12,   {"b1": null, "b2": "4321"}, "bar" ]
    ])");

  // 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), record_batch->schema(),
                             arrow_writer_properties, &arrow_writer));
  // NewBufferedRowGroup() is not called explicitly and it will be called
  // inside WriteRecordBatch().
  ASSERT_OK_NO_THROW(arrow_writer->WriteRecordBatch(*record_batch));
  ASSERT_OK_NO_THROW(arrow_writer->Close());
  ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish());

  // Create reader with batch size specified.
  auto read_properties = default_arrow_reader_properties();
  read_properties.set_batch_size(record_batch->num_rows());
  auto reader = ParquetFileReader::Open(std::make_shared<BufferReader>(buffer));
  std::unique_ptr<FileReader> arrow_reader;
  ASSERT_OK_AND_ASSIGN(arrow_reader,
                       FileReader::Make(pool, std::move(reader), read_properties));

  // Verify the single record batch has been sliced into two row groups by
  // WriterProperties::max_row_group_length().
  auto file_metadata = arrow_reader->parquet_reader()->metadata();
  int num_row_groups = file_metadata->num_row_groups();
  ASSERT_EQ(2, num_row_groups);
  ASSERT_EQ(10, file_metadata->RowGroup(0)->num_rows());
  ASSERT_EQ(2, file_metadata->RowGroup(1)->num_rows());

  // Verify that page index is not written.
  for (int i = 0; i < num_row_groups; ++i) {
    auto row_group_metadata = file_metadata->RowGroup(i);
    for (int j = 0; j < row_group_metadata->num_columns(); ++j) {
      auto column_metadata = row_group_metadata->ColumnChunk(j);
      EXPECT_FALSE(column_metadata->GetColumnIndexLocation().has_value());
      EXPECT_FALSE(column_metadata->GetOffsetIndexLocation().has_value());
    }
  }

  // Verify batch data read via RecordBatch
  ASSERT_OK_AND_ASSIGN(auto batch_reader,
                       arrow_reader->GetRecordBatchReader(Iota(num_row_groups)));
  std::shared_ptr<::arrow::RecordBatch> read_record_batch;
  ASSERT_OK(batch_reader->ReadNext(&read_record_batch));

Review comment (suggestion):
```suggestion
  ASSERT_OK_AND_ASSIGN(auto arrow_reader,
```

Code after:

  // 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), record_batch->schema(),
                             arrow_writer_properties, &arrow_writer));
  // NewBufferedRowGroup() is not called explicitly and it will be called
  // inside WriteRecordBatch().
  ASSERT_OK_NO_THROW(arrow_writer->WriteRecordBatch(*record_batch));
  ASSERT_OK_NO_THROW(arrow_writer->Close());
  ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish());

  // Create reader with batch size specified.
  auto read_properties = default_arrow_reader_properties();
  read_properties.set_batch_size(record_batch->num_rows());
  auto reader = ParquetFileReader::Open(std::make_shared<BufferReader>(buffer));
  ASSERT_OK_AND_ASSIGN(auto arrow_reader,
                       FileReader::Make(pool, std::move(reader), read_properties));

  // Verify the single record batch has been sliced into two row groups by
  // WriterProperties::max_row_group_length().
  auto file_metadata = arrow_reader->parquet_reader()->metadata();
  int num_row_groups = file_metadata->num_row_groups();
  ASSERT_EQ(2, num_row_groups);
  ASSERT_EQ(10, file_metadata->RowGroup(0)->num_rows());
  ASSERT_EQ(2, file_metadata->RowGroup(1)->num_rows());

  // Verify that page index is not written.
  for (int i = 0; i < num_row_groups; ++i) {
    auto row_group_metadata = file_metadata->RowGroup(i);
    for (int j = 0; j < row_group_metadata->num_columns(); ++j) {
      auto column_metadata = row_group_metadata->ColumnChunk(j);
      EXPECT_FALSE(column_metadata->GetColumnIndexLocation().has_value());
      EXPECT_FALSE(column_metadata->GetOffsetIndexLocation().has_value());
    }
  }

  // Verify batch data read via RecordBatch
  ASSERT_OK_AND_ASSIGN(auto batch_reader,
                       arrow_reader->GetRecordBatchReader(Iota(num_row_groups)));
  std::shared_ptr<::arrow::RecordBatch> read_record_batch;
  ASSERT_OK(batch_reader->ReadNext(&read_record_batch));
  EXPECT_TRUE(record_batch->Equals(*read_record_batch));
}

TEST(TestArrowReadWrite, WriteRecordBatchNotProduceEmptyRowGroup) {
  // GH-39211: WriteRecordBatch should prevent from writing a empty row group
  // in the end of the file.
  auto pool = ::arrow::default_memory_pool();
  auto sink = CreateOutputStream();
  // Limit the max number of rows in a row group to 2