From 725fe70db6efcd45bbdbdd6f3ef97c8c71221fad Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Fri, 7 Aug 2026 10:30:29 -0400 Subject: [PATCH] Minor: migrate more arrow round trip tests to use RoundTripTest fixture --- parquet/src/arrow/arrow_writer/mod.rs | 127 ++++++++------------------ 1 file changed, 38 insertions(+), 89 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 0b4cdfff42ac..4ffdb17c964d 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -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(); } #[test] @@ -2261,15 +2257,11 @@ mod tests { .build() .unwrap(); let a = ListArray::from(a_list_data); + assert_eq!(a.null_count(), 1); - // 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] @@ -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] @@ -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] @@ -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] @@ -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] @@ -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] @@ -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(), @@ -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(); + 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(), @@ -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(); + 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. @@ -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] { - 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 {