Skip to content

Minor: migrate more arrow round trip tests to use RoundTripTest fixture #10585

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
alamb merged 1 commit into from
Aug 9, 2026
Merged

Minor: migrate more arrow round trip tests to use RoundTripTest fixture #10585

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
127 changes: 38 additions & 89 deletions parquet/src/arrow/arrow_writer/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2221,16 +2221,12 @@ mod tests {

#[test]
fn arrow_writer_non_null() {
// define schema
let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);

// create some data
let a = Int32Array::from(vec![1, 2, 3, 4, 5]);

// build a record batch
let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();

roundtrip(batch, Some(SMALL_SIZE / 2));
RoundTripTest::new(Arc::new(a))
.with_schema(Arc::new(schema))
.run();

@alamb alamb Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This actually increases test coverage as

RoundTripTest::run() actually already calls into the same roundtrip code, but does so with:

  1. Multiple different row group sizes (including SMALL_SIZE/2)
  2. All supported encodings
  3. Dictionary/no dictionary pages
  4. DataPage V1 and V2

https://github.com/apache/arrow-rs/blob/725fe70db6efcd45bbdbdd6f3ef97c8c71221fad/parquet/src/arrow/arrow_writer/mod.rs#L3208-L3207

}

#[test]
Expand Down Expand Up @@ -2261,15 +2257,11 @@ mod tests {
.build()
.unwrap();
let a = ListArray::from(a_list_data);
assert_eq!(a.null_count(), 1);

@alamb alamb Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

as above, this actually increases coverage substantially


// build a record batch
let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();

assert_eq!(batch.column(0).null_count(), 1);

// This test fails if the max row group size is less than the batch's length
// see https://github.com/apache/arrow-rs/issues/518
roundtrip(batch, None);
RoundTripTest::new(Arc::new(a))
.with_schema(Arc::new(schema))
.run();
}

#[test]
Expand Down Expand Up @@ -2299,15 +2291,11 @@ mod tests {
.build()
.unwrap();
let a = ListArray::from(a_list_data);
assert_eq!(a.null_count(), 0);

// build a record batch
let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();

// This test fails if the max row group size is less than the batch's length
// see https://github.com/apache/arrow-rs/issues/518
assert_eq!(batch.column(0).null_count(), 0);

roundtrip(batch, None);
RoundTripTest::new(Arc::new(a))
.with_schema(Arc::new(schema))
.run();
}

#[test]
Expand All @@ -2327,12 +2315,11 @@ mod tests {
Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
Some(vec![true, true, false, true, true].into()),
);
assert_eq!(a.null_count(), 1);

let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();

assert_eq!(batch.column(0).null_count(), 1);

roundtrip(batch, None);
RoundTripTest::new(Arc::new(a))
.with_schema(Arc::new(schema))
.run();
}

#[test]
Expand All @@ -2352,12 +2339,11 @@ mod tests {
Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
None,
);
assert_eq!(a.null_count(), 0);

let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();

assert_eq!(batch.column(0).null_count(), 0);

roundtrip(batch, None);
RoundTripTest::new(Arc::new(a))
.with_schema(Arc::new(schema))
.run();
}

#[test]
Expand All @@ -2377,10 +2363,11 @@ mod tests {
Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
None,
);
assert_eq!(a.null_count(), 0);

let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();

roundtrip(batch, None);
RoundTripTest::new(Arc::new(a))
.with_schema(Arc::new(schema))
.run();
}

#[test]
Expand All @@ -2400,12 +2387,11 @@ mod tests {
Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
Some(vec![true, true, false, true, true].into()),
);
assert_eq!(a.null_count(), 1);

let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();

assert_eq!(batch.column(0).null_count(), 1);

roundtrip(batch, None);
RoundTripTest::new(Arc::new(a))
.with_schema(Arc::new(schema))
.run();
}

#[test]
Expand Down Expand Up @@ -2441,18 +2427,15 @@ mod tests {
Arc::new(struct_array),
Some(vec![true, false, true].into()),
);
assert_eq!(list_view.null_count(), 1);

let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(list_view)]).unwrap();

roundtrip(batch, None);
RoundTripTest::new(Arc::new(list_view))
.with_schema(Arc::new(schema))
.run();
}

#[test]
fn arrow_writer_binary() {
let string_field = Field::new("a", DataType::Utf8, false);
let binary_field = Field::new("b", DataType::Binary, false);
let schema = Schema::new(vec![string_field, binary_field]);

let raw_string_values = vec!["foo", "bar", "baz", "quux"];
let raw_binary_values = [
b"foo".to_vec(),
Expand All @@ -2467,22 +2450,15 @@ mod tests {

let string_values = StringArray::from(raw_string_values.clone());
let binary_values = BinaryArray::from(raw_binary_value_refs);
let batch = RecordBatch::try_new(
Arc::new(schema),
vec![Arc::new(string_values), Arc::new(binary_values)],
)
.unwrap();
assert_eq!(string_values.null_count(), 0);
assert_eq!(binary_values.null_count(), 0);

roundtrip(batch, Some(SMALL_SIZE / 2));
RoundTripTest::new(Arc::new(string_values)).run();

@alamb alamb Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

technically speaking this encodes a single column batch rather than a multi-column batch, but the point of these tests is to test the round tripping of data through the whole parquet machiner; Batches with multiple columns are tested elsewhere

RoundTripTest::new(Arc::new(binary_values)).run();
}

#[test]
fn arrow_writer_binary_view() {
let string_field = Field::new("a", DataType::Utf8View, false);
let binary_field = Field::new("b", DataType::BinaryView, false);
let nullable_string_field = Field::new("a", DataType::Utf8View, true);
let schema = Schema::new(vec![string_field, binary_field, nullable_string_field]);

let raw_string_values = vec!["foo", "bar", "large payload over 12 bytes", "lulu"];
let raw_binary_values = vec![
b"foo".to_vec(),
Expand All @@ -2496,26 +2472,14 @@ mod tests {
let string_view_values = StringViewArray::from(raw_string_values);
let binary_view_values = BinaryViewArray::from_iter_values(raw_binary_values);
let nullable_string_view_values = StringViewArray::from(nullable_string_values);
let batch = RecordBatch::try_new(
Arc::new(schema),
vec![
Arc::new(string_view_values),
Arc::new(binary_view_values),
Arc::new(nullable_string_view_values),
],
)
.unwrap();

roundtrip(batch.clone(), Some(SMALL_SIZE / 2));
roundtrip(batch, None);
RoundTripTest::new(Arc::new(string_view_values)).run();

@alamb alamb Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As above, this ctually increases coverage because it tries a both these two sets of max row group sizes, and a bunch of other permutations

RoundTripTest::new(Arc::new(binary_view_values)).run();
RoundTripTest::new(Arc::new(nullable_string_view_values)).run();
}

#[test]
fn arrow_writer_binary_view_long_value() {
let string_field = Field::new("a", DataType::Utf8View, false);
let binary_field = Field::new("b", DataType::BinaryView, false);
let schema = Schema::new(vec![string_field, binary_field]);

// There is special case validation for long values (greater than 128)
// 128 encodes as 0x80 0x00 0x00 0x00 in little endian, which should
// trigger the long-string UTF-8 validation branch in the plain decoder.
Expand All @@ -2533,21 +2497,6 @@ mod tests {
RoundTripTest::new(Arc::clone(&binary_view_values))
.with_nullable(false)
.run();

let batch = RecordBatch::try_new(
Arc::new(schema),
vec![string_view_values, binary_view_values],
)
.unwrap();

// Disable dictionary to exercise plain encoding paths in the reader.
for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {

@alamb alamb Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

RoundtripTest also tests with/without dictionary and data page formats (and several other parameters)

https://github.com/apache/arrow-rs/blob/725fe70db6efcd45bbdbdd6f3ef97c8c71221fad/parquet/src/arrow/arrow_writer/mod.rs#L3199-L3198

let props = WriterProperties::builder()
.set_writer_version(version)
.set_dictionary_enabled(false)
.build();
roundtrip_opts(&batch, props);
}
}

fn get_decimal_batch(precision: u8, scale: i8) -> RecordBatch {
Expand Down
Loading