This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 88b7d5ab9a Encapsulate parquet round trip test logic in struct (#10545)
88b7d5ab9a is described below
commit 88b7d5ab9a43796fd115b5f502fb408c15a6f3af
Author: Andrew Lamb <[email protected]>
AuthorDate: Thu Aug 6 16:23:15 2026 -0400
Encapsulate parquet round trip test logic in struct (#10545)
# Which issue does this PR close?
- Part of https://github.com/apache/arrow-rs/issues/10540
# Rationale for this change
I am trying to understand the depth of our round trip tests (there are
over 4000
lines of tests in the arrow writer module). After some study it appears
there
are tests to round trip both single columns and record batches, which
share some
non trivial amount of logic.
It also makes it hard to evaluate coverage because there are several
similar but not quite the same free functions and it is unclear hwo they
are related to each other
and what some of the parameters mean (like the argument to
`roundtrip_one_column`)
I think making it easier to find and evaluate test coverage will make it
easier to
maintain and extend this crate in the future.
# What changes are included in this PR?
1. Move round trip logic into methods on `RoundTripTest` rather than
free functions
# Are these changes tested?
Only tests
# Are there any user-facing changes?
No
---------
Co-authored-by: Claude Fable 5 <[email protected]>
---
parquet/src/arrow/arrow_writer/mod.rs | 298 ++++++++++++++++++++--------------
1 file changed, 172 insertions(+), 126 deletions(-)
diff --git a/parquet/src/arrow/arrow_writer/mod.rs
b/parquet/src/arrow/arrow_writer/mod.rs
index 337ec60184..6c3ce7811b 100644
--- a/parquet/src/arrow/arrow_writer/mod.rs
+++ b/parquet/src/arrow/arrow_writer/mod.rs
@@ -2527,8 +2527,12 @@ mod tests {
let binary_view_values: ArrayRef =
Arc::new(BinaryViewArray::from_iter_values(raw_binary_values));
- one_column_roundtrip(Arc::clone(&string_view_values), false);
- one_column_roundtrip(Arc::clone(&binary_view_values), false);
+ RoundTripTest::new(Arc::clone(&string_view_values))
+ .with_nullable(false)
+ .run();
+ RoundTripTest::new(Arc::clone(&binary_view_values))
+ .with_nullable(false)
+ .run();
let batch = RecordBatch::try_new(
Arc::new(schema),
@@ -3127,101 +3131,138 @@ mod tests {
})
}
- struct RoundTripOptions {
+ /// Round trip testing fixture:
+ ///
+ /// Tests based on this fixture write data to parquet and then read it
back.
+ struct RoundTripTest {
values: ArrayRef,
- schema: SchemaRef,
+ /// Optionally supplied schema
+ schema: Option<SchemaRef>,
+ /// If the created schema should be nullable. Defaults to true. Ignored
+ /// if schema is set to Some.
+ nullable: bool,
bloom_filter: bool,
bloom_filter_ndv: Option<u64>,
bloom_filter_position: BloomFilterPosition,
}
- impl RoundTripOptions {
- fn new(values: ArrayRef, nullable: bool) -> Self {
- let data_type = values.data_type().clone();
- let schema = Schema::new(vec![Field::new("col", data_type,
nullable)]);
+ impl RoundTripTest {
+ /// Create a test for round tripping values with a nullable schema
+ fn new(values: ArrayRef) -> Self {
Self {
values,
- schema: Arc::new(schema),
+ schema: None,
+ nullable: true,
bloom_filter: false,
bloom_filter_ndv: None,
bloom_filter_position: BloomFilterPosition::AfterRowGroup,
}
}
- }
- fn one_column_roundtrip(values: ArrayRef, nullable: bool) -> Vec<Bytes> {
- one_column_roundtrip_with_options(RoundTripOptions::new(values,
nullable))
- }
+ /// Set the schema
+ fn with_schema(mut self, schema: SchemaRef) -> Self {
+ self.schema = Some(schema);
+ self
+ }
- fn one_column_roundtrip_with_schema(values: ArrayRef, schema: SchemaRef)
-> Vec<Bytes> {
- let mut options = RoundTripOptions::new(values, false);
- options.schema = schema;
- one_column_roundtrip_with_options(options)
- }
+ /// Set the nullable flag
+ fn with_nullable(mut self, nullable: bool) -> Self {
+ self.nullable = nullable;
+ self
+ }
- fn one_column_roundtrip_with_options(options: RoundTripOptions) ->
Vec<Bytes> {
- let RoundTripOptions {
- values,
- schema,
- bloom_filter,
- bloom_filter_ndv,
- bloom_filter_position,
- } = options;
-
- let encodings = match values.data_type() {
- DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary |
DataType::LargeBinary => {
- vec![
- Encoding::PLAIN,
- Encoding::DELTA_BYTE_ARRAY,
- Encoding::DELTA_LENGTH_BYTE_ARRAY,
- ]
- }
- DataType::Int64
- | DataType::Int32
- | DataType::Int16
- | DataType::Int8
- | DataType::UInt64
- | DataType::UInt32
- | DataType::UInt16
- | DataType::UInt8 => vec![
- Encoding::PLAIN,
- Encoding::DELTA_BINARY_PACKED,
- Encoding::BYTE_STREAM_SPLIT,
- ],
- DataType::Float32 | DataType::Float64 => {
- vec![Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT]
- }
- _ => vec![Encoding::PLAIN],
- };
+ /// Set bloom filter
+ fn with_bloom_filter(mut self, bloom_filter: bool) -> Self {
+ self.bloom_filter = bloom_filter;
+ self
+ }
- let expected_batch = RecordBatch::try_new(schema,
vec![values]).unwrap();
+ /// Set bloom filter max ndv
+ fn with_bloom_filter_ndv(mut self, bloom_filter_ndv: u64) -> Self {
+ self.bloom_filter_ndv = Some(bloom_filter_ndv);
+ self
+ }
- let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE /
2 + 1, 10];
+ /// Set bloom filter position
+ fn with_bloom_filter_position(
+ mut self,
+ bloom_filter_position: BloomFilterPosition,
+ ) -> Self {
+ self.bloom_filter_position = bloom_filter_position;
+ self
+ }
- let mut files = vec![];
- for dictionary_size in [0, 1, 1024] {
- for encoding in &encodings {
- for version in [WriterVersion::PARQUET_1_0,
WriterVersion::PARQUET_2_0] {
- for row_group_size in row_group_sizes {
- let mut builder = WriterProperties::builder()
- .set_writer_version(version)
- .set_max_row_group_row_count(Some(row_group_size))
- .set_dictionary_enabled(dictionary_size != 0)
-
.set_dictionary_page_size_limit(dictionary_size.max(1))
- .set_encoding(*encoding)
- .set_bloom_filter_enabled(bloom_filter)
- .set_bloom_filter_position(bloom_filter_position);
- if let Some(ndv) = bloom_filter_ndv {
- builder = builder.set_bloom_filter_max_ndv(ndv);
- }
- let props = builder.build();
+ /// Run the test specified by the options, returning the encoded
Parquet bytes
+ fn run(self) -> Vec<Bytes> {
+ let RoundTripTest {
+ values,
+ schema,
+ nullable,
+ bloom_filter,
+ bloom_filter_ndv,
+ bloom_filter_position,
+ } = self;
+
+ let schema = schema.unwrap_or_else(|| {
+ let data_type = values.data_type().clone();
+ Arc::new(Schema::new(vec![Field::new("col", data_type,
nullable)]))
+ });
+
+ let encodings = match values.data_type() {
+ DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary |
DataType::LargeBinary => {
+ vec![
+ Encoding::PLAIN,
+ Encoding::DELTA_BYTE_ARRAY,
+ Encoding::DELTA_LENGTH_BYTE_ARRAY,
+ ]
+ }
+ DataType::Int64
+ | DataType::Int32
+ | DataType::Int16
+ | DataType::Int8
+ | DataType::UInt64
+ | DataType::UInt32
+ | DataType::UInt16
+ | DataType::UInt8 => vec![
+ Encoding::PLAIN,
+ Encoding::DELTA_BINARY_PACKED,
+ Encoding::BYTE_STREAM_SPLIT,
+ ],
+ DataType::Float32 | DataType::Float64 => {
+ vec![Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT]
+ }
+ _ => vec![Encoding::PLAIN],
+ };
- files.push(roundtrip_opts(&expected_batch, props))
+ let expected_batch = RecordBatch::try_new(schema,
vec![values]).unwrap();
+
+ let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2,
SMALL_SIZE / 2 + 1, 10];
+
+ let mut files = vec![];
+ for dictionary_size in [0, 1, 1024] {
+ for encoding in &encodings {
+ for version in [WriterVersion::PARQUET_1_0,
WriterVersion::PARQUET_2_0] {
+ for row_group_size in row_group_sizes {
+ let mut builder = WriterProperties::builder()
+ .set_writer_version(version)
+
.set_max_row_group_row_count(Some(row_group_size))
+ .set_dictionary_enabled(dictionary_size != 0)
+
.set_dictionary_page_size_limit(dictionary_size.max(1))
+ .set_encoding(*encoding)
+ .set_bloom_filter_enabled(bloom_filter)
+
.set_bloom_filter_position(bloom_filter_position);
+ if let Some(ndv) = bloom_filter_ndv {
+ builder =
builder.set_bloom_filter_max_ndv(ndv);
+ }
+ let props = builder.build();
+
+ files.push(roundtrip_opts(&expected_batch, props))
+ }
}
}
}
+ files
}
- files
}
fn values_required<A, I>(iter: I) -> Vec<Bytes>
@@ -3231,7 +3272,7 @@ mod tests {
{
let raw_values: Vec<_> = iter.into_iter().collect();
let values = Arc::new(A::from(raw_values));
- one_column_roundtrip(values, false)
+ RoundTripTest::new(values).with_nullable(false).run()
}
fn values_optional<A, I>(iter: I) -> Vec<Bytes>
@@ -3245,7 +3286,7 @@ mod tests {
.map(|(i, v)| if i % 2 == 0 { None } else { Some(v) })
.collect();
let optional_values = Arc::new(A::from(optional_raw_values));
- one_column_roundtrip(optional_values, true)
+ RoundTripTest::new(optional_values).run()
}
fn required_and_optional<A, I>(iter: I)
@@ -3322,12 +3363,12 @@ mod tests {
#[test]
fn all_null_primitive_single_column() {
let values = Arc::new(Int32Array::from(vec![None; SMALL_SIZE]));
- one_column_roundtrip(values, true);
+ RoundTripTest::new(values).run();
}
#[test]
fn null_single_column() {
let values = Arc::new(NullArray::new(SMALL_SIZE));
- one_column_roundtrip(values, true);
+ RoundTripTest::new(values).run();
// null arrays are always nullable, a test with non-nullable nulls
fails
}
@@ -3542,14 +3583,14 @@ mod tests {
// The timestamp array types don't implement From<Vec<T>> because they
need the timezone
// argument, and they also doesn't support building from a Vec<Option<T>>,
so call
- // one_column_roundtrip manually instead of calling required_and_optional
for these tests.
+ // RoundTripTest manually instead of calling required_and_optional for
these tests.
#[test]
fn timestamp_second_single_column() {
let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
let values = Arc::new(TimestampSecondArray::from(raw_values));
- one_column_roundtrip(values, false);
+ RoundTripTest::new(values).with_nullable(false).run();
}
#[test]
@@ -3557,7 +3598,7 @@ mod tests {
let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
let values = Arc::new(TimestampMillisecondArray::from(raw_values));
- one_column_roundtrip(values, false);
+ RoundTripTest::new(values).with_nullable(false).run();
}
#[test]
@@ -3565,7 +3606,7 @@ mod tests {
let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
let values = Arc::new(TimestampMicrosecondArray::from(raw_values));
- one_column_roundtrip(values, false);
+ RoundTripTest::new(values).with_nullable(false).run();
}
#[test]
@@ -3573,7 +3614,7 @@ mod tests {
let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
let values = Arc::new(TimestampNanosecondArray::from(raw_values));
- one_column_roundtrip(values, false);
+ RoundTripTest::new(values).with_nullable(false).run();
}
#[test]
@@ -3680,11 +3721,12 @@ mod tests {
#[test]
fn i32_column_bloom_filter_at_end() {
let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
- let mut options = RoundTripOptions::new(array, false);
- options.bloom_filter = true;
- options.bloom_filter_position = BloomFilterPosition::End;
+ let files = RoundTripTest::new(array)
+ .with_nullable(false)
+ .with_bloom_filter(true)
+ .with_bloom_filter_position(BloomFilterPosition::End)
+ .run();
- let files = one_column_roundtrip_with_options(options);
check_bloom_filter(
files,
"col".to_string(),
@@ -3696,10 +3738,11 @@ mod tests {
#[test]
fn i32_column_bloom_filter() {
let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
- let mut options = RoundTripOptions::new(array, false);
- options.bloom_filter = true;
+ let files = RoundTripTest::new(array)
+ .with_nullable(false)
+ .with_bloom_filter(true)
+ .run();
- let files = one_column_roundtrip_with_options(options);
check_bloom_filter(
files,
"col".to_string(),
@@ -3717,11 +3760,12 @@ mod tests {
let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
// NDV much larger than actual distinct values — tests folding a large
filter down
- let mut options = RoundTripOptions::new(array.clone(), false);
- options.bloom_filter = true;
- options.bloom_filter_ndv = Some(1_000_000);
+ let files = RoundTripTest::new(array.clone())
+ .with_nullable(false)
+ .with_bloom_filter(true)
+ .with_bloom_filter_ndv(1_000_000)
+ .run();
- let files = one_column_roundtrip_with_options(options);
check_bloom_filter(
files,
"col".to_string(),
@@ -3730,11 +3774,12 @@ mod tests {
);
// NDV smaller than actual distinct values — tests the underestimate
path
- let mut options = RoundTripOptions::new(array, false);
- options.bloom_filter = true;
- options.bloom_filter_ndv = Some(3);
+ let files = RoundTripTest::new(array)
+ .with_nullable(false)
+ .with_bloom_filter(true)
+ .with_bloom_filter_ndv(3)
+ .run();
- let files = one_column_roundtrip_with_options(options);
check_bloom_filter(
files,
"col".to_string(),
@@ -3750,10 +3795,11 @@ mod tests {
let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
let array = Arc::new(BinaryArray::from_iter_values(many_vecs_iter));
- let mut options = RoundTripOptions::new(array, false);
- options.bloom_filter = true;
+ let files = RoundTripTest::new(array)
+ .with_nullable(false)
+ .with_bloom_filter(true)
+ .run();
- let files = one_column_roundtrip_with_options(options);
check_bloom_filter(
files,
"col".to_string(),
@@ -3768,10 +3814,10 @@ mod tests {
let raw_strs = raw_values.iter().map(|s| s.as_str());
let array = Arc::new(StringArray::from_iter_values(raw_strs));
- let mut options = RoundTripOptions::new(array, false);
- options.bloom_filter = true;
-
- let files = one_column_roundtrip_with_options(options);
+ let files = RoundTripTest::new(array)
+ .with_nullable(false)
+ .with_bloom_filter(true)
+ .run();
let optional_raw_values: Vec<_> = raw_values
.iter()
@@ -3801,7 +3847,7 @@ mod tests {
builder.append_value(b"1112").unwrap();
let array = Arc::new(builder.finish());
- one_column_roundtrip(array, true);
+ RoundTripTest::new(array).run();
}
#[test]
@@ -3883,7 +3929,7 @@ mod tests {
let a = ListArray::from(a_list_data);
let values = Arc::new(a);
- one_column_roundtrip(values, true);
+ RoundTripTest::new(values).run();
}
#[test]
@@ -3908,7 +3954,7 @@ mod tests {
let a = LargeListArray::from(a_list_data);
let values = Arc::new(a);
- one_column_roundtrip(values, true);
+ RoundTripTest::new(values).run();
}
#[test]
@@ -3924,10 +3970,10 @@ mod tests {
];
let list = ListArray::from_iter_primitive::<Int32Type, _,
_>(data.clone());
- one_column_roundtrip(Arc::new(list), true);
+ RoundTripTest::new(Arc::new(list)).run();
let list = LargeListArray::from_iter_primitive::<Int32Type, _,
_>(data);
- one_column_roundtrip(Arc::new(list), true);
+ RoundTripTest::new(Arc::new(list)).run();
}
#[test]
@@ -3944,7 +3990,7 @@ mod tests {
builder.values().append_value("large payload over 12 bytes");
builder.append(true);
- one_column_roundtrip(Arc::new(builder.finish()), true);
+ RoundTripTest::new(Arc::new(builder.finish())).run();
}
#[test]
@@ -3954,7 +4000,7 @@ mod tests {
let s = StructArray::from(vec![(struct_field_a, Arc::new(a_values) as
ArrayRef)]);
let values = Arc::new(s);
- one_column_roundtrip(values, false);
+ RoundTripTest::new(values).with_nullable(false).run();
}
#[test]
@@ -4072,7 +4118,7 @@ mod tests {
.collect();
// build a record batch
- one_column_roundtrip_with_schema(Arc::new(d), schema);
+ RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
}
#[test]
@@ -4326,7 +4372,7 @@ mod tests {
builder.append(12345678).unwrap();
let d = builder.finish();
- one_column_roundtrip_with_schema(Arc::new(d), schema);
+ RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
}
#[test]
@@ -4340,14 +4386,14 @@ mod tests {
.unwrap();
let array = DictionaryArray::new(keys, Arc::new(values));
- one_column_roundtrip(Arc::new(array.clone()), true);
+ RoundTripTest::new(Arc::new(array.clone())).run();
let values = Decimal32Array::from(integers)
.with_precision_and_scale(9, 2)
.unwrap();
let array = array.with_values(Arc::new(values));
- one_column_roundtrip(Arc::new(array), true);
+ RoundTripTest::new(Arc::new(array)).run();
}
#[test]
@@ -4361,14 +4407,14 @@ mod tests {
.unwrap();
let array = DictionaryArray::new(keys, Arc::new(values));
- one_column_roundtrip(Arc::new(array.clone()), true);
+ RoundTripTest::new(Arc::new(array.clone())).run();
let values = Decimal64Array::from(integers)
.with_precision_and_scale(12, 2)
.unwrap();
let array = array.with_values(Arc::new(values));
- one_column_roundtrip(Arc::new(array), true);
+ RoundTripTest::new(Arc::new(array)).run();
}
#[test]
@@ -4382,14 +4428,14 @@ mod tests {
.unwrap();
let array = DictionaryArray::new(keys, Arc::new(values));
- one_column_roundtrip(Arc::new(array.clone()), true);
+ RoundTripTest::new(Arc::new(array.clone())).run();
let values = Decimal128Array::from(integers)
.with_precision_and_scale(12, 2)
.unwrap();
let array = array.with_values(Arc::new(values));
- one_column_roundtrip(Arc::new(array), true);
+ RoundTripTest::new(Arc::new(array)).run();
}
#[test]
@@ -4407,14 +4453,14 @@ mod tests {
.unwrap();
let array = DictionaryArray::new(keys, Arc::new(values));
- one_column_roundtrip(Arc::new(array.clone()), true);
+ RoundTripTest::new(Arc::new(array.clone())).run();
let values = Decimal256Array::from(integers)
.with_precision_and_scale(12, 2)
.unwrap();
let array = array.with_values(Arc::new(values));
- one_column_roundtrip(Arc::new(array), true);
+ RoundTripTest::new(Arc::new(array)).run();
}
#[test]
@@ -4435,7 +4481,7 @@ mod tests {
.copied()
.collect();
- one_column_roundtrip_with_schema(Arc::new(d), schema);
+ RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
}
#[test]
@@ -4451,7 +4497,7 @@ mod tests {
u32::MAX,
];
let values =
Arc::new(UInt32Array::from_iter_values(src.iter().cloned()));
- let files = one_column_roundtrip(values, false);
+ let files = RoundTripTest::new(values).with_nullable(false).run();
for file in files {
// check statistics are valid
@@ -4497,7 +4543,7 @@ mod tests {
u64::MAX,
];
let values =
Arc::new(UInt64Array::from_iter_values(src.iter().cloned()));
- let files = one_column_roundtrip(values, false);
+ let files = RoundTripTest::new(values).with_nullable(false).run();
for file in files {
// check statistics are valid
@@ -4534,7 +4580,7 @@ mod tests {
fn statistics_null_counts_only_nulls() {
// check that null-count statistics for "only NULL"-columns are correct
let values = Arc::new(UInt64Array::from(vec![None, None]));
- let files = one_column_roundtrip(values, true);
+ let files = RoundTripTest::new(values).run();
for file in files {
// check statistics are valid
@@ -4636,7 +4682,7 @@ mod tests {
let array = Arc::new(list_builder.finish());
- one_column_roundtrip(array, true);
+ RoundTripTest::new(array).run();
}
fn row_group_sizes(metadata: &ParquetMetaData) -> Vec<i64> {
