[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