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 05bcf41997 test(parquet): move scalar roundtrip tests (#11128)
05bcf41997 is described below
commit 05bcf41997c5c350a6ab9e1d816c898f9d9fb24f
Author: Phoenix <[email protected]>
AuthorDate: Fri Sep 25 19:37:47 2026 +0800
test(parquet): move scalar roundtrip tests (#11128)
# Which issue does this PR close?
Part of #10540.
# Rationale for this change
Continue consolidating Arrow/Parquet roundtrip coverage in the writer
integration suite so related end-to-end tests are easier to find and
extend. Sharing the existing harness allows the tests to move in small
batches while the remaining unit tests keep using the same validation
logic.
# What changes are included in this PR?
- Move 31 tests from the writer's unit test module into
`parquet/tests/arrow_writer/roundtrip.rs`: 10 numeric, 14 temporal, and
7 binary/string tests.
- Extract the existing roundtrip helpers into a test-only module shared
by unit and integration tests. No public library API or feature is
added.
- Preserve the test bodies, assertions,
encoding/dictionary/writer-version/row-group combinations, and Miri
annotations. The integration module temporarily expects unused helpers
because only part of the suite has moved.
# Are these changes tested?
- `cargo test --locked --offline -p parquet --lib
arrow::arrow_writer::tests::`: 106 passed.
- `cargo test --locked --offline -p parquet --test arrow_writer`: 58
passed.
- `cargo test --locked --offline -p parquet --no-default-features
--features arrow --test arrow_writer`: 58 passed.
- `cargo fmt --all --check`
- `cargo clippy --locked --offline -p parquet --all-targets
--all-features -- -D warnings`
The affected writer suites still contain 164 tests in total: 137 + 27
before migration, 106 + 58 after migration.
# Are there any user-facing changes?
No. This PR reorganizes tests only.
---
parquet/src/arrow/arrow_writer/mod.rs | 480 +-----------------------
parquet/tests/arrow_writer/mod.rs | 5 +
parquet/tests/arrow_writer/roundtrip.rs | 250 +++++++++++-
parquet/tests/arrow_writer/roundtrip_helpers.rs | 267 +++++++++++++
4 files changed, 530 insertions(+), 472 deletions(-)
diff --git a/parquet/src/arrow/arrow_writer/mod.rs
b/parquet/src/arrow/arrow_writer/mod.rs
index 29b413531c..0c61143948 100644
--- a/parquet/src/arrow/arrow_writer/mod.rs
+++ b/parquet/src/arrow/arrow_writer/mod.rs
@@ -2075,8 +2075,20 @@ fn update_distinct_values_seen(
}
}
+// Allow the helpers to use the same imports in unit and integration tests.
+#[cfg(test)]
+use crate as parquet_crate;
+
+#[cfg(test)]
+#[path = "../../../tests/arrow_writer/roundtrip_helpers.rs"]
+mod roundtrip_helpers;
+
#[cfg(test)]
mod tests {
+ use super::roundtrip_helpers::{
+ RoundTripTest, SMALL_SIZE, required_and_optional, roundtrip,
roundtrip_opts,
+ roundtrip_opts_with_array_validation,
+ };
use super::*;
use std::cmp::Ordering;
use std::collections::HashMap;
@@ -3237,240 +3249,8 @@ mod tests {
roundtrip(batch, None);
}
- const SMALL_SIZE: usize = 7;
const MEDIUM_SIZE: usize = 63;
- // Write the batch to parquet and read it back out, ensuring
- // that what comes out is the same as what was written in
- fn roundtrip(expected_batch: RecordBatch, max_row_group_size:
Option<usize>) -> Vec<Bytes> {
- let mut files = vec![];
- for version in [WriterVersion::PARQUET_1_0,
WriterVersion::PARQUET_2_0] {
- let mut props =
WriterProperties::builder().set_writer_version(version);
-
- if let Some(size) = max_row_group_size {
- props = props.set_max_row_group_row_count(Some(size))
- }
-
- let props = props.build();
- files.push(roundtrip_opts(&expected_batch, props))
- }
- files
- }
-
- // Round trip the specified record batch with the specified writer
properties,
- // to an in-memory file, and validate the arrays using the specified
function.
- // Returns the in-memory file.
- fn roundtrip_opts_with_array_validation<F>(
- expected_batch: &RecordBatch,
- props: WriterProperties,
- validate: F,
- ) -> Bytes
- where
- F: Fn(&ArrayData, &ArrayData),
- {
- let mut file = vec![];
-
- let mut writer = ArrowWriter::try_new(&mut file,
expected_batch.schema(), Some(props))
- .expect("Unable to write file");
- writer.write(expected_batch).unwrap();
- writer.close().unwrap();
-
- let file = Bytes::from(file);
- let mut record_batch_reader =
- ParquetRecordBatchReader::try_new(file.clone(), 1024).unwrap();
-
- let actual_batch = record_batch_reader
- .next()
- .expect("No batch found")
- .expect("Unable to get batch");
-
- assert_eq!(expected_batch.schema(), actual_batch.schema());
- assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
- assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
- for i in 0..expected_batch.num_columns() {
- let expected_data = expected_batch.column(i).to_data();
- let actual_data = actual_batch.column(i).to_data();
- validate(&expected_data, &actual_data);
- }
-
- file
- }
-
- fn roundtrip_opts(expected_batch: &RecordBatch, props: WriterProperties)
-> Bytes {
- roundtrip_opts_with_array_validation(expected_batch, props, |a, b| {
- a.validate_full().expect("valid expected data");
- b.validate_full().expect("valid actual data");
- assert_eq!(a, b)
- })
- }
-
- /// Round trip testing fixture:
- ///
- /// Tests based on this fixture write data to parquet and then read it
back.
- struct RoundTripTest {
- values: ArrayRef,
- /// 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 RoundTripTest {
- /// Create a test for round tripping values with a nullable schema
- fn new(values: ArrayRef) -> Self {
- Self {
- values,
- schema: None,
- nullable: true,
- bloom_filter: false,
- bloom_filter_ndv: None,
- bloom_filter_position: BloomFilterPosition::AfterRowGroup,
- }
- }
-
- /// Set the schema
- fn with_schema(mut self, schema: SchemaRef) -> Self {
- self.schema = Some(schema);
- self
- }
-
- /// Set the nullable flag
- fn with_nullable(mut self, nullable: bool) -> Self {
- self.nullable = nullable;
- self
- }
-
- /// Set bloom filter
- fn with_bloom_filter(mut self, bloom_filter: bool) -> Self {
- self.bloom_filter = bloom_filter;
- self
- }
-
- /// 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
- }
-
- /// Set bloom filter position
- fn with_bloom_filter_position(
- mut self,
- bloom_filter_position: BloomFilterPosition,
- ) -> Self {
- self.bloom_filter_position = bloom_filter_position;
- self
- }
-
- /// 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,
Encoding::ALP]
- }
- _ => vec![Encoding::PLAIN],
- };
-
- 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
- }
- }
-
- fn values_required<A, I>(iter: I) -> Vec<Bytes>
- where
- A: From<Vec<I::Item>> + Array + 'static,
- I: IntoIterator,
- {
- let raw_values: Vec<_> = iter.into_iter().collect();
- let values = Arc::new(A::from(raw_values));
- RoundTripTest::new(values).with_nullable(false).run()
- }
-
- fn values_optional<A, I>(iter: I) -> Vec<Bytes>
- where
- A: From<Vec<Option<I::Item>>> + Array + 'static,
- I: IntoIterator,
- {
- let optional_raw_values: Vec<_> = iter
- .into_iter()
- .enumerate()
- .map(|(i, v)| if i % 2 == 0 { None } else { Some(v) })
- .collect();
- let optional_values = Arc::new(A::from(optional_raw_values));
- RoundTripTest::new(optional_values).run()
- }
-
- fn required_and_optional<A, I>(iter: I)
- where
- A: From<Vec<I::Item>> + From<Vec<Option<I::Item>>> + Array + 'static,
- I: IntoIterator + Clone,
- {
- values_required::<A, I>(iter.clone());
- values_optional::<A, I>(iter);
- }
-
fn check_bloom_filter<T: AsBytes>(
files: Vec<Bytes>,
file_column: String,
@@ -3710,169 +3490,6 @@ mod tests {
assert_eq!(float_idx.max_value(3), Some(&1.0));
}
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn i8_single_column() {
- required_and_optional::<Int8Array, _>(0..SMALL_SIZE as i8);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn i16_single_column() {
- required_and_optional::<Int16Array, _>(0..SMALL_SIZE as i16);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn i32_single_column() {
- required_and_optional::<Int32Array, _>(0..SMALL_SIZE as i32);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn i64_single_column() {
- required_and_optional::<Int64Array, _>(0..SMALL_SIZE as i64);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn u8_single_column() {
- required_and_optional::<UInt8Array, _>(0..SMALL_SIZE as u8);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn u16_single_column() {
- required_and_optional::<UInt16Array, _>(0..SMALL_SIZE as u16);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn u32_single_column() {
- required_and_optional::<UInt32Array, _>(0..SMALL_SIZE as u32);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn u64_single_column() {
- required_and_optional::<UInt64Array, _>(0..SMALL_SIZE as u64);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn f32_single_column() {
- required_and_optional::<Float32Array, _>((0..SMALL_SIZE).map(|i| i as
f32));
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn f64_single_column() {
- required_and_optional::<Float64Array, _>((0..SMALL_SIZE).map(|i| i as
f64));
- }
-
- // 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
- // RoundTripTest manually instead of calling required_and_optional for
these tests.
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn timestamp_second_single_column() {
- let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
- let values = Arc::new(TimestampSecondArray::from(raw_values));
-
- RoundTripTest::new(values).with_nullable(false).run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn timestamp_millisecond_single_column() {
- let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
- let values = Arc::new(TimestampMillisecondArray::from(raw_values));
-
- RoundTripTest::new(values).with_nullable(false).run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn timestamp_microsecond_single_column() {
- let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
- let values = Arc::new(TimestampMicrosecondArray::from(raw_values));
-
- RoundTripTest::new(values).with_nullable(false).run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn timestamp_nanosecond_single_column() {
- let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
- let values = Arc::new(TimestampNanosecondArray::from(raw_values));
-
- RoundTripTest::new(values).with_nullable(false).run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn date32_single_column() {
- required_and_optional::<Date32Array, _>(0..SMALL_SIZE as i32);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn date64_single_column() {
- // Date64 must be a multiple of 86400000, see ARROW-10925
- required_and_optional::<Date64Array, _>(
- (0..(SMALL_SIZE as i64 * 86400000)).step_by(86400000),
- );
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn time32_second_single_column() {
- required_and_optional::<Time32SecondArray, _>(0..SMALL_SIZE as i32);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn time32_millisecond_single_column() {
- required_and_optional::<Time32MillisecondArray, _>(0..SMALL_SIZE as
i32);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn time64_microsecond_single_column() {
- required_and_optional::<Time64MicrosecondArray, _>(0..SMALL_SIZE as
i64);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn time64_nanosecond_single_column() {
- required_and_optional::<Time64NanosecondArray, _>(0..SMALL_SIZE as
i64);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn duration_second_single_column() {
- required_and_optional::<DurationSecondArray, _>(0..SMALL_SIZE as i64);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn duration_millisecond_single_column() {
- required_and_optional::<DurationMillisecondArray, _>(0..SMALL_SIZE as
i64);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn duration_microsecond_single_column() {
- required_and_optional::<DurationMicrosecondArray, _>(0..SMALL_SIZE as
i64);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn duration_nanosecond_single_column() {
- required_and_optional::<DurationNanosecondArray, _>(0..SMALL_SIZE as
i64);
- }
-
#[test]
#[cfg_attr(miri, ignore)] // Takes too long
fn interval_year_month_single_column() {
@@ -3903,28 +3520,6 @@ mod tests {
]);
}
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn binary_single_column() {
- let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
- let many_vecs: Vec<_> = std::iter::repeat_n(one_vec,
SMALL_SIZE).collect();
- let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
-
- // BinaryArrays can't be built from Vec<Option<&str>>, so only call
`values_required`
- values_required::<BinaryArray, _>(many_vecs_iter);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn binary_view_single_column() {
- let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
- let many_vecs: Vec<_> = std::iter::repeat_n(one_vec,
SMALL_SIZE).collect();
- let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
-
- // BinaryArrays can't be built from Vec<Option<&str>>, so only call
`values_required`
- values_required::<BinaryViewArray, _>(many_vecs_iter);
- }
-
#[test]
#[cfg_attr(miri, ignore)] // Takes too long
fn i32_column_bloom_filter_at_end() {
@@ -4143,57 +3738,6 @@ mod tests {
check_bloom_filter(files, "col".to_string(), optional_raw_values,
vec![""]);
}
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn large_binary_single_column() {
- let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
- let many_vecs: Vec<_> = std::iter::repeat_n(one_vec,
SMALL_SIZE).collect();
- let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
-
- // LargeBinaryArrays can't be built from Vec<Option<&str>>, so only
call `values_required`
- values_required::<LargeBinaryArray, _>(many_vecs_iter);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn fixed_size_binary_single_column() {
- let mut builder = FixedSizeBinaryBuilder::new(4);
- builder.append_value(b"0123").unwrap();
- builder.append_null();
- builder.append_value(b"8910").unwrap();
- builder.append_value(b"1112").unwrap();
- let array = Arc::new(builder.finish());
-
- RoundTripTest::new(array).run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn string_single_column() {
- let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i|
i.to_string()).collect();
- let raw_strs = raw_values.iter().map(|s| s.as_str());
-
- required_and_optional::<StringArray, _>(raw_strs);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn large_string_single_column() {
- let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i|
i.to_string()).collect();
- let raw_strs = raw_values.iter().map(|s| s.as_str());
-
- required_and_optional::<LargeStringArray, _>(raw_strs);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn string_view_single_column() {
- let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i|
i.to_string()).collect();
- let raw_strs = raw_values.iter().map(|s| s.as_str());
-
- required_and_optional::<StringViewArray, _>(raw_strs);
- }
-
#[test]
fn null_list_single_column() {
let null_field = Field::new_list_field(DataType::Null, true);
diff --git a/parquet/tests/arrow_writer/mod.rs
b/parquet/tests/arrow_writer/mod.rs
index a6e243eb0d..195e855c40 100644
--- a/parquet/tests/arrow_writer/mod.rs
+++ b/parquet/tests/arrow_writer/mod.rs
@@ -19,6 +19,11 @@
mod layout;
mod roundtrip;
+// The integration tests currently use only part of the shared helpers.
+#[expect(dead_code)]
+mod roundtrip_helpers;
+
+use parquet as parquet_crate;
use std::alloc::{GlobalAlloc, Layout, System};
use std::cell::Cell;
diff --git a/parquet/tests/arrow_writer/roundtrip.rs
b/parquet/tests/arrow_writer/roundtrip.rs
index dd4b9ec4d4..135dd0df77 100644
--- a/parquet/tests/arrow_writer/roundtrip.rs
+++ b/parquet/tests/arrow_writer/roundtrip.rs
@@ -17,19 +17,27 @@
//! Round-trip tests for Arrow data written to Parquet.
+use super::roundtrip_helpers::{RoundTripTest, SMALL_SIZE,
required_and_optional, values_required};
+
use std::collections::HashMap;
use std::sync::Arc;
+use arrow_array::builder::FixedSizeBinaryBuilder;
use arrow_array::cast::AsArray;
use arrow_array::types::{
Date32Type, Date64Type, Decimal32Type, Decimal64Type, Decimal128Type,
Decimal256Type,
DecimalType, Float16Type, Time32MillisecondType, Time64MicrosecondType,
};
use arrow_array::{
- Array, ArrayRef, Decimal128Array, Decimal256Array, DictionaryArray,
FixedSizeBinaryArray,
- Float16Array, Int32Array, ListArray, PrimitiveArray, RecordBatch,
RecordBatchReader,
- StringArray, StructArray, Time32MillisecondArray, Time64MicrosecondArray,
UInt8Array,
- UInt8DictionaryArray, UInt32Array, UInt64Array,
+ Array, ArrayRef, BinaryArray, BinaryViewArray, Date32Array, Date64Array,
Decimal128Array,
+ Decimal256Array, DictionaryArray, DurationMicrosecondArray,
DurationMillisecondArray,
+ DurationNanosecondArray, DurationSecondArray, FixedSizeBinaryArray,
Float16Array, Float32Array,
+ Float64Array, Int8Array, Int16Array, Int32Array, Int64Array,
LargeBinaryArray,
+ LargeStringArray, ListArray, PrimitiveArray, RecordBatch,
RecordBatchReader, StringArray,
+ StringViewArray, StructArray, Time32MillisecondArray, Time32SecondArray,
+ Time64MicrosecondArray, Time64NanosecondArray, TimestampMicrosecondArray,
+ TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray,
UInt8Array,
+ UInt8DictionaryArray, UInt16Array, UInt32Array, UInt64Array,
};
use arrow_buffer::{ArrowNativeType, Buffer, NullBuffer, i256};
use arrow_data::ArrayDataBuilder;
@@ -43,6 +51,240 @@ use parquet::basic::Type as PhysicalType;
use parquet::errors::Result;
use parquet::file::properties::WriterProperties;
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn i8_single_column() {
+ required_and_optional::<Int8Array, _>(0..SMALL_SIZE as i8);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn i16_single_column() {
+ required_and_optional::<Int16Array, _>(0..SMALL_SIZE as i16);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn i32_single_column() {
+ required_and_optional::<Int32Array, _>(0..SMALL_SIZE as i32);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn i64_single_column() {
+ required_and_optional::<Int64Array, _>(0..SMALL_SIZE as i64);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn u8_single_column() {
+ required_and_optional::<UInt8Array, _>(0..SMALL_SIZE as u8);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn u16_single_column() {
+ required_and_optional::<UInt16Array, _>(0..SMALL_SIZE as u16);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn u32_single_column() {
+ required_and_optional::<UInt32Array, _>(0..SMALL_SIZE as u32);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn u64_single_column() {
+ required_and_optional::<UInt64Array, _>(0..SMALL_SIZE as u64);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn f32_single_column() {
+ required_and_optional::<Float32Array, _>((0..SMALL_SIZE).map(|i| i as
f32));
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn f64_single_column() {
+ required_and_optional::<Float64Array, _>((0..SMALL_SIZE).map(|i| i as
f64));
+}
+
+// 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
+// RoundTripTest manually instead of calling required_and_optional for these
tests.
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn timestamp_second_single_column() {
+ let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
+ let values = Arc::new(TimestampSecondArray::from(raw_values));
+
+ RoundTripTest::new(values).with_nullable(false).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn timestamp_millisecond_single_column() {
+ let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
+ let values = Arc::new(TimestampMillisecondArray::from(raw_values));
+
+ RoundTripTest::new(values).with_nullable(false).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn timestamp_microsecond_single_column() {
+ let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
+ let values = Arc::new(TimestampMicrosecondArray::from(raw_values));
+
+ RoundTripTest::new(values).with_nullable(false).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn timestamp_nanosecond_single_column() {
+ let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
+ let values = Arc::new(TimestampNanosecondArray::from(raw_values));
+
+ RoundTripTest::new(values).with_nullable(false).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn date32_single_column() {
+ required_and_optional::<Date32Array, _>(0..SMALL_SIZE as i32);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn date64_single_column() {
+ // Date64 must be a multiple of 86400000, see ARROW-10925
+ required_and_optional::<Date64Array, _>((0..(SMALL_SIZE as i64 *
86400000)).step_by(86400000));
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn time32_second_single_column() {
+ required_and_optional::<Time32SecondArray, _>(0..SMALL_SIZE as i32);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn time32_millisecond_single_column() {
+ required_and_optional::<Time32MillisecondArray, _>(0..SMALL_SIZE as i32);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn time64_microsecond_single_column() {
+ required_and_optional::<Time64MicrosecondArray, _>(0..SMALL_SIZE as i64);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn time64_nanosecond_single_column() {
+ required_and_optional::<Time64NanosecondArray, _>(0..SMALL_SIZE as i64);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn duration_second_single_column() {
+ required_and_optional::<DurationSecondArray, _>(0..SMALL_SIZE as i64);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn duration_millisecond_single_column() {
+ required_and_optional::<DurationMillisecondArray, _>(0..SMALL_SIZE as i64);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn duration_microsecond_single_column() {
+ required_and_optional::<DurationMicrosecondArray, _>(0..SMALL_SIZE as i64);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn duration_nanosecond_single_column() {
+ required_and_optional::<DurationNanosecondArray, _>(0..SMALL_SIZE as i64);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn binary_single_column() {
+ let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
+ let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
+ let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
+
+ // BinaryArrays can't be built from Vec<Option<&str>>, so only call
`values_required`
+ values_required::<BinaryArray, _>(many_vecs_iter);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn binary_view_single_column() {
+ let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
+ let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
+ let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
+
+ // BinaryArrays can't be built from Vec<Option<&str>>, so only call
`values_required`
+ values_required::<BinaryViewArray, _>(many_vecs_iter);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn large_binary_single_column() {
+ let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
+ let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
+ let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
+
+ // LargeBinaryArrays can't be built from Vec<Option<&str>>, so only call
`values_required`
+ values_required::<LargeBinaryArray, _>(many_vecs_iter);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn fixed_size_binary_single_column() {
+ let mut builder = FixedSizeBinaryBuilder::new(4);
+ builder.append_value(b"0123").unwrap();
+ builder.append_null();
+ builder.append_value(b"8910").unwrap();
+ builder.append_value(b"1112").unwrap();
+ let array = Arc::new(builder.finish());
+
+ RoundTripTest::new(array).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn string_single_column() {
+ let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
+ let raw_strs = raw_values.iter().map(|s| s.as_str());
+
+ required_and_optional::<StringArray, _>(raw_strs);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn large_string_single_column() {
+ let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
+ let raw_strs = raw_values.iter().map(|s| s.as_str());
+
+ required_and_optional::<LargeStringArray, _>(raw_strs);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn string_view_single_column() {
+ let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
+ let raw_strs = raw_values.iter().map(|s| s.as_str());
+
+ required_and_optional::<StringViewArray, _>(raw_strs);
+}
+
#[test]
fn test_unsigned_roundtrip() {
let schema = Arc::new(Schema::new(vec![
diff --git a/parquet/tests/arrow_writer/roundtrip_helpers.rs
b/parquet/tests/arrow_writer/roundtrip_helpers.rs
new file mode 100644
index 0000000000..5d541e227a
--- /dev/null
+++ b/parquet/tests/arrow_writer/roundtrip_helpers.rs
@@ -0,0 +1,267 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Shared round-trip helpers for Arrow writer tests.
+
+use std::sync::Arc;
+
+use arrow_array::{Array, ArrayRef, RecordBatch};
+use arrow_data::ArrayData;
+use arrow_schema::{DataType, Field, Schema, SchemaRef};
+use bytes::Bytes;
+
+use super::parquet_crate::arrow::ArrowWriter;
+use super::parquet_crate::arrow::arrow_reader::ParquetRecordBatchReader;
+use super::parquet_crate::basic::Encoding;
+use super::parquet_crate::file::properties::{
+ BloomFilterPosition, WriterProperties, WriterVersion,
+};
+
+pub(super) const SMALL_SIZE: usize = 7;
+
+// Write the batch to parquet and read it back out, ensuring
+// that what comes out is the same as what was written in
+pub(super) fn roundtrip(
+ expected_batch: RecordBatch,
+ max_row_group_size: Option<usize>,
+) -> Vec<Bytes> {
+ let mut files = vec![];
+ for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
+ let mut props =
WriterProperties::builder().set_writer_version(version);
+
+ if let Some(size) = max_row_group_size {
+ props = props.set_max_row_group_row_count(Some(size))
+ }
+
+ let props = props.build();
+ files.push(roundtrip_opts(&expected_batch, props))
+ }
+ files
+}
+
+// Round trip the specified record batch with the specified writer properties,
+// to an in-memory file, and validate the arrays using the specified function.
+// Returns the in-memory file.
+pub(super) fn roundtrip_opts_with_array_validation<F>(
+ expected_batch: &RecordBatch,
+ props: WriterProperties,
+ validate: F,
+) -> Bytes
+where
+ F: Fn(&ArrayData, &ArrayData),
+{
+ let mut file = vec![];
+
+ let mut writer = ArrowWriter::try_new(&mut file, expected_batch.schema(),
Some(props))
+ .expect("Unable to write file");
+ writer.write(expected_batch).unwrap();
+ writer.close().unwrap();
+
+ let file = Bytes::from(file);
+ let mut record_batch_reader =
ParquetRecordBatchReader::try_new(file.clone(), 1024).unwrap();
+
+ let actual_batch = record_batch_reader
+ .next()
+ .expect("No batch found")
+ .expect("Unable to get batch");
+
+ assert_eq!(expected_batch.schema(), actual_batch.schema());
+ assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
+ assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
+ for i in 0..expected_batch.num_columns() {
+ let expected_data = expected_batch.column(i).to_data();
+ let actual_data = actual_batch.column(i).to_data();
+ validate(&expected_data, &actual_data);
+ }
+
+ file
+}
+
+pub(super) fn roundtrip_opts(expected_batch: &RecordBatch, props:
WriterProperties) -> Bytes {
+ roundtrip_opts_with_array_validation(expected_batch, props, |a, b| {
+ a.validate_full().expect("valid expected data");
+ b.validate_full().expect("valid actual data");
+ assert_eq!(a, b)
+ })
+}
+
+/// Round trip testing fixture:
+///
+/// Tests based on this fixture write data to parquet and then read it back.
+pub(super) struct RoundTripTest {
+ values: ArrayRef,
+ /// 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 RoundTripTest {
+ /// Create a test for round tripping values with a nullable schema
+ pub(super) fn new(values: ArrayRef) -> Self {
+ Self {
+ values,
+ schema: None,
+ nullable: true,
+ bloom_filter: false,
+ bloom_filter_ndv: None,
+ bloom_filter_position: BloomFilterPosition::AfterRowGroup,
+ }
+ }
+
+ /// Set the schema
+ pub(super) fn with_schema(mut self, schema: SchemaRef) -> Self {
+ self.schema = Some(schema);
+ self
+ }
+
+ /// Set the nullable flag
+ pub(super) fn with_nullable(mut self, nullable: bool) -> Self {
+ self.nullable = nullable;
+ self
+ }
+
+ /// Set bloom filter
+ pub(super) fn with_bloom_filter(mut self, bloom_filter: bool) -> Self {
+ self.bloom_filter = bloom_filter;
+ self
+ }
+
+ /// Set bloom filter max ndv
+ pub(super) fn with_bloom_filter_ndv(mut self, bloom_filter_ndv: u64) ->
Self {
+ self.bloom_filter_ndv = Some(bloom_filter_ndv);
+ self
+ }
+
+ /// Set bloom filter position
+ pub(super) fn with_bloom_filter_position(
+ mut self,
+ bloom_filter_position: BloomFilterPosition,
+ ) -> Self {
+ self.bloom_filter_position = bloom_filter_position;
+ self
+ }
+
+ /// Run the test specified by the options, returning the encoded Parquet
bytes
+ pub(super) 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,
Encoding::ALP]
+ }
+ _ => vec![Encoding::PLAIN],
+ };
+
+ 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
+ }
+}
+
+pub(super) fn values_required<A, I>(iter: I) -> Vec<Bytes>
+where
+ A: From<Vec<I::Item>> + Array + 'static,
+ I: IntoIterator,
+{
+ let raw_values: Vec<_> = iter.into_iter().collect();
+ let values = Arc::new(A::from(raw_values));
+ RoundTripTest::new(values).with_nullable(false).run()
+}
+
+fn values_optional<A, I>(iter: I) -> Vec<Bytes>
+where
+ A: From<Vec<Option<I::Item>>> + Array + 'static,
+ I: IntoIterator,
+{
+ let optional_raw_values: Vec<_> = iter
+ .into_iter()
+ .enumerate()
+ .map(|(i, v)| if i % 2 == 0 { None } else { Some(v) })
+ .collect();
+ let optional_values = Arc::new(A::from(optional_raw_values));
+ RoundTripTest::new(optional_values).run()
+}
+
+pub(super) fn required_and_optional<A, I>(iter: I)
+where
+ A: From<Vec<I::Item>> + From<Vec<Option<I::Item>>> + Array + 'static,
+ I: IntoIterator + Clone,
+{
+ values_required::<A, I>(iter.clone());
+ values_optional::<A, I>(iter);
+}
