[C++][Parquet] Page sizes for repeated columns can overflow int32 with page index enabled

[adamreeve]

Describe the bug, including details regarding any error messages, version, and platform.

I noticed a regression when upgrading from Arrow 19.0.1 to 20.0.0 and writing Parquet files with a repeated column. It appears that the change to enable writing the page index by default (#45249) has caused the logic for starting new data pages to change, and so page sizes can become very large, and can overflow int32.

Repro code, from commit adamreeve@8444bf6:

TEST(TestColumnWriter, WriteLargeLists) {
  auto sink = CreateOutputStream();
  auto schema = std::static_pointer_cast<GroupNode>(GroupNode::Make(
      "schema", Repetition::REQUIRED,
      {
          GroupNode::Make(
              "x", Repetition::OPTIONAL,
              {
                  GroupNode::Make("list", Repetition::REPEATED,
                                  {
                                      schema::Float("element", Repetition::REQUIRED),
                                  },
                                  nullptr),
              },
              LogicalType::List()),
      }));
  auto properties = WriterProperties::Builder()
                        .disable_dictionary()
                        //->disable_write_page_index()
                        ->build();
  auto file_writer = ParquetFileWriter::Open(sink, schema, properties);
  auto rg_writer = file_writer->AppendRowGroup();

  constexpr int64_t num_rows = 1000 * 1000;
  constexpr int64_t num_list_elements = 1000;

  std::vector<int16_t> def_levels(num_list_elements);
  std::vector<int16_t> rep_levels(num_list_elements);
  std::vector<float> values(num_list_elements);

  for (int32_t i = 0; i < num_list_elements; ++i) {
    def_levels[i] = 2;
    rep_levels[i] = i == 0 ? 0 : 1;
  }

  auto col_writer = dynamic_cast<FloatWriter*>(rg_writer->NextColumn());
  for (int32_t i = 0; i < num_rows; i++) {
    random_numbers(num_list_elements, i, -100.0f, 100.0f, values.data());
    col_writer->WriteBatch(num_list_elements, def_levels.data(), rep_levels.data(),
                           values.data());
  }
  col_writer->Close();

  rg_writer->Close();
  file_writer->Close();
  ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish());
}

This runs with disable_write_page_index(), but otherwise crashes with:

C++ exception with description "Uncompressed data page size overflows INT32_MAX. Size:4005000014" thrown in the test body.

Component(s)

C++, Parquet

[adamreeve]

This seems to be triggered by writing batches aligned with whole lists, which isn't handled by the logic for finding a record boundary in DoInBatches:

arrow/cpp/src/parquet/column_writer.cc

Lines 1175 to 1209 in b712c1e

	int64_t offset = 0;
	while (offset < num_levels) {
	int64_t end_offset = std::min(offset + batch_size, num_levels);
	// Find next record boundary (i.e. ref_level = 0)
	while (end_offset < num_levels && rep_levels[end_offset] != 0) {
	end_offset++;
	}
	if (end_offset < num_levels) {
	// This is not the last chunk of batch and end_offset is a record boundary.
	// It is a good chance to check the page size.
	action(offset, end_offset - offset, /*check_page_size=*/true);
	} else {
	DCHECK_EQ(end_offset, num_levels);
	// This is the last chunk of batch, and we do not know whether end_offset is a
	// record boundary. Find the offset to beginning of last record in this chunk,
	// so we can check page size.
	int64_t last_record_begin_offset = num_levels - 1;
	while (last_record_begin_offset >= offset &&
	rep_levels[last_record_begin_offset] != 0) {
	last_record_begin_offset--;
	}
	if (offset < last_record_begin_offset) {
	// We have found the beginning of last record and can check page size.
	action(offset, last_record_begin_offset - offset, /*check_page_size=*/true);
	offset = last_record_begin_offset;
	}
	// There is no record boundary in this chunk and cannot check page size.
	action(offset, end_offset - offset, /*check_page_size=*/false);
	}
	offset = end_offset;

I think the fix might just be:

--- a/cpp/src/parquet/column_writer.cc
+++ b/cpp/src/parquet/column_writer.cc
@@ -1196,7 +1196,7 @@ inline void DoInBatches(const int16_t* def_levels, const int16_t* rep_levels,
         last_record_begin_offset--;
       }
 
-      if (offset < last_record_begin_offset) {
+      if (offset <= last_record_begin_offset) {
         // We have found the beginning of last record and can check page size.
         action(offset, last_record_begin_offset - offset, /*check_page_size=*/true);
         offset = last_record_begin_offset;

[pitrou]

I think the fix might just be:

I think that's right indeed.

[adamreeve]

Issue resolved by pull request 47032
#47032