You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by gi...@apache.org on 2026/10/06 04:41:24 UTC
(arrow-rs) 01/01: test(parquet): move REE and basic roundtrips —— Part III (#11184)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11184-788040d30e4dd2cc245a76717ded95a5e313475f
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit c5e390031a1443f97355ed80411d6885e14642ff
Author: Phoenix <Ph...@163.com>
AuthorDate: Tue Oct 6 04:41:04 2026 +0000
test(parquet): move REE and basic roundtrips —— Part III (#11184)
# Which issue does this PR close?
Part of apache/arrow-rs#10540.
# Stack base
This draft PR targets `apache/arrow-rs:main`. It follows stage II, which
was merged in apache/arrow-rs#11183. The branch now contains only the
two stage-III commits on top of `main`, so this PR's diff is the
stage-III change.
# Rationale for this change
Continue moving writer roundtrip coverage in reviewable batches.
# What changes are included in this PR?
Move run-end-encoded (REE) and remaining basic writer roundtrip tests
into `parquet/tests/arrow_writer/roundtrip.rs`. The two commits preserve
their data comparisons and test options.
# Are these changes tested?
- `cargo test --locked -p parquet --test arrow_writer` — 107 passed.
- `cargo test --locked -p parquet --lib -- arrow::arrow_writer::tests` —
61 passed.
- `cargo fmt --all -- --check`
- `git diff --check upstream/main..HEAD`
# Are there any user-facing changes?
No. This is a test-only move.
---
parquet/src/arrow/arrow_writer/mod.rs | 389 +---------------------
parquet/tests/arrow_writer/mod.rs | 2 -
parquet/tests/arrow_writer/roundtrip.rs | 415 +++++++++++++++++++++++-
parquet/tests/arrow_writer/roundtrip_helpers.rs | 27 +-
4 files changed, 426 insertions(+), 407 deletions(-)
diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs
index a72584bbd8..0bfb1c4f5e 100644
--- a/parquet/src/arrow/arrow_writer/mod.rs
+++ b/parquet/src/arrow/arrow_writer/mod.rs
@@ -2082,7 +2082,7 @@ mod roundtrip_helpers;
#[cfg(test)]
mod tests {
use super::roundtrip_helpers::{
- RoundTripTest, SMALL_SIZE, required_and_optional, roundtrip, roundtrip_opts,
+ RoundTripTest, SMALL_SIZE, required_and_optional, roundtrip_opts,
roundtrip_opts_with_array_validation,
};
use super::*;
@@ -2104,9 +2104,8 @@ mod tests {
use arrow::util::data_gen::create_random_array;
use arrow::util::pretty::pretty_format_batches;
use arrow::{array::*, buffer::Buffer};
- use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, NullBuffer, OffsetBuffer};
+ use arrow_buffer::{IntervalMonthDayNano, NullBuffer, OffsetBuffer};
use arrow_schema::Fields;
- use half::f16;
use tempfile::tempfile;
use crate::basic::{Encoding, EncodingMask};
@@ -2333,24 +2332,6 @@ mod tests {
);
}
- #[test]
- fn arrow_writer() {
- // define schema
- let schema = Schema::new(vec![
- Field::new("a", DataType::Int32, false),
- Field::new("b", DataType::Int32, true),
- ]);
-
- // create some data
- let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
- let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
-
- // build a record batch
- let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
-
- roundtrip(batch, Some(SMALL_SIZE / 2));
- }
-
fn get_bytes_after_close(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
let mut buffer = vec![];
@@ -2407,112 +2388,6 @@ mod tests {
}
}
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_non_null() {
- let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
- let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
-
- RoundTripTest::new(Arc::new(a))
- .with_schema(Arc::new(schema))
- .run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_binary() {
- let raw_string_values = vec!["foo", "bar", "baz", "quux"];
- let raw_binary_values = [
- b"foo".to_vec(),
- b"bar".to_vec(),
- b"baz".to_vec(),
- b"quux".to_vec(),
- ];
- let raw_binary_value_refs = raw_binary_values
- .iter()
- .map(|x| x.as_slice())
- .collect::<Vec<_>>();
-
- let string_values = StringArray::from(raw_string_values.clone());
- let binary_values = BinaryArray::from(raw_binary_value_refs);
- assert_eq!(string_values.null_count(), 0);
- assert_eq!(binary_values.null_count(), 0);
-
- RoundTripTest::new(Arc::new(string_values)).run();
- RoundTripTest::new(Arc::new(binary_values)).run();
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_binary_view() {
- let raw_string_values = vec!["foo", "bar", "large payload over 12 bytes", "lulu"];
- let raw_binary_values = vec![
- b"foo".to_vec(),
- b"bar".to_vec(),
- b"large payload over 12 bytes".to_vec(),
- b"lulu".to_vec(),
- ];
- let nullable_string_values =
- vec![Some("foo"), None, Some("large payload over 12 bytes"), None];
-
- 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);
-
- 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]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn arrow_writer_binary_view_long_value() {
- // 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.
- let long = "a".repeat(128);
- let raw_string_values = vec!["foo", long.as_str(), "bar"];
- let raw_binary_values = vec![b"foo".to_vec(), long.as_bytes().to_vec(), b"bar".to_vec()];
-
- let string_view_values: ArrayRef = Arc::new(StringViewArray::from(raw_string_values));
- let binary_view_values: ArrayRef =
- Arc::new(BinaryViewArray::from_iter_values(raw_binary_values));
-
- RoundTripTest::new(Arc::clone(&string_view_values))
- .with_nullable(false)
- .run();
- RoundTripTest::new(Arc::clone(&binary_view_values))
- .with_nullable(false)
- .run();
- }
-
- fn get_decimal_batch(precision: u8, scale: i8) -> RecordBatch {
- let decimal_field = Field::new("a", DataType::Decimal128(precision, scale), false);
- let schema = Schema::new(vec![decimal_field]);
-
- let decimal_values = vec![10_000, 50_000, 0, -100]
- .into_iter()
- .map(Some)
- .collect::<Decimal128Array>()
- .with_precision_and_scale(precision, scale)
- .unwrap();
-
- RecordBatch::try_new(Arc::new(schema), vec![Arc::new(decimal_values)]).unwrap()
- }
-
- #[test]
- fn arrow_writer_decimal() {
- // int32 to store the decimal value
- let batch_int32_decimal = get_decimal_batch(5, 2);
- roundtrip(batch_int32_decimal, Some(SMALL_SIZE / 2));
- // int64 to store the decimal value
- let batch_int64_decimal = get_decimal_batch(12, 2);
- roundtrip(batch_int64_decimal, Some(SMALL_SIZE / 2));
- // fixed_length_byte_array to store the decimal value
- let batch_fixed_len_byte_array_decimal = get_decimal_batch(30, 2);
- roundtrip(batch_fixed_len_byte_array_decimal, Some(SMALL_SIZE / 2));
- }
-
#[test]
fn arrow_writer_page_size() {
let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
@@ -2582,45 +2457,6 @@ mod tests {
);
}
- #[test]
- #[cfg_attr(miri, ignore)] // inline assembly is not supported
- fn arrow_writer_float_nans() {
- let f16_field = Field::new("a", DataType::Float16, false);
- let f32_field = Field::new("b", DataType::Float32, false);
- let f64_field = Field::new("c", DataType::Float64, false);
- let schema = Schema::new(vec![f16_field, f32_field, f64_field]);
-
- let f16_values = (0..MEDIUM_SIZE)
- .map(|i| {
- Some(if i % 2 == 0 {
- f16::NAN
- } else {
- f16::from_f32(i as f32)
- })
- })
- .collect::<Float16Array>();
-
- let f32_values = (0..MEDIUM_SIZE)
- .map(|i| Some(if i % 2 == 0 { f32::NAN } else { i as f32 }))
- .collect::<Float32Array>();
-
- let f64_values = (0..MEDIUM_SIZE)
- .map(|i| Some(if i % 2 == 0 { f64::NAN } else { i as f64 }))
- .collect::<Float64Array>();
-
- let batch = RecordBatch::try_new(
- Arc::new(schema),
- vec![
- Arc::new(f16_values),
- Arc::new(f32_values),
- Arc::new(f64_values),
- ],
- )
- .unwrap();
-
- roundtrip(batch, None);
- }
-
const MEDIUM_SIZE: usize = 63;
fn check_bloom_filter<T: AsBytes>(
@@ -2685,28 +2521,6 @@ mod tests {
});
}
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn all_null_primitive_single_column() {
- let values = Arc::new(Int32Array::from(vec![None; SMALL_SIZE]));
- RoundTripTest::new(values).run();
- }
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn null_single_column() {
- let values = Arc::new(NullArray::new(SMALL_SIZE));
- RoundTripTest::new(values).run();
- // null arrays are always nullable, a test with non-nullable nulls fails
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn bool_single_column() {
- required_and_optional::<BooleanArray, _>(
- [true, false].iter().cycle().copied().take(SMALL_SIZE),
- );
- }
-
#[test]
#[cfg_attr(miri, ignore)] // Takes too long
fn bool_large_single_column() {
@@ -2862,23 +2676,6 @@ mod tests {
assert_eq!(float_idx.max_value(3), Some(&1.0));
}
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn interval_year_month_single_column() {
- required_and_optional::<IntervalYearMonthArray, _>(0..SMALL_SIZE as i32);
- }
-
- #[test]
- #[cfg_attr(miri, ignore)] // Takes too long
- fn interval_day_time_single_column() {
- required_and_optional::<IntervalDayTimeArray, _>(vec![
- IntervalDayTime::new(0, 1),
- IntervalDayTime::new(0, 3),
- IntervalDayTime::new(3, -2),
- IntervalDayTime::new(-200, 4),
- ]);
- }
-
#[test]
#[should_panic(
expected = "Attempting to write an Arrow interval type MonthDayNano to parquet that is not yet implemented"
@@ -4896,188 +4693,6 @@ mod tests {
assert!(cc.column_index_range().is_none());
}
- /// Writes a single-column RecordBatch to an in-memory Parquet buffer.
- fn write_column_to_bytes(array: ArrayRef) -> Bytes {
- let schema = Arc::new(Schema::new(vec![Field::new(
- "col",
- array.data_type().clone(),
- true,
- )]));
- let buf = get_bytes_after_close(
- schema.clone(),
- &RecordBatch::try_new(schema, vec![array]).unwrap(),
- );
- Bytes::from(buf)
- }
-
- /// Reads column 0 from a single-row-group Parquet buffer, projecting it with the given schema.
- /// Passing a flat schema when the buffer was written from a REE array lets callers decode
- /// the physical values without the run-end encoding wrapper.
- fn read_column_with_schema(bytes: Bytes, schema: SchemaRef) -> ArrayRef {
- let opts = crate::arrow::arrow_reader::ArrowReaderOptions::new().with_schema(schema);
- ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, opts)
- .unwrap()
- .build()
- .unwrap()
- .next()
- .unwrap()
- .unwrap()
- .column(0)
- .clone()
- }
-
- fn ree_write_read_roundtrip(ree: ArrayRef, flat: ArrayRef) {
- let flat_schema = Arc::new(Schema::new(vec![Field::new(
- "col",
- flat.data_type().clone(),
- true,
- )]));
- let ree_bytes = write_column_to_bytes(ree);
- let flat_bytes = write_column_to_bytes(flat.clone());
- assert_eq!(
- ree_bytes, flat_bytes,
- "REE and flat bytes should be identical"
- );
-
- let decoded_ree = read_column_with_schema(ree_bytes, flat_schema.clone());
- let decoded_flat = read_column_with_schema(flat_bytes, flat_schema);
-
- assert_eq!(decoded_ree.as_ref(), flat.as_ref());
- assert_eq!(decoded_ree.as_ref(), decoded_flat.as_ref());
- }
-
- #[test]
- fn ree_string() {
- let ree: ArrayRef = Arc::new(
- [Some("a"), Some("a"), None, Some("b"), Some("b")]
- .into_iter()
- .collect::<Int32RunArray>(),
- );
- let flat: ArrayRef = Arc::new(StringArray::from(vec![
- Some("a"),
- Some("a"),
- None,
- Some("b"),
- Some("b"),
- ]));
- ree_write_read_roundtrip(ree, flat);
- }
-
- #[test]
- fn ree_int32() {
- let mut b = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
- for v in [Some(1), Some(1), None, Some(2), Some(2)] {
- b.append_option(v);
- }
- let ree: ArrayRef = Arc::new(b.finish());
- let flat: ArrayRef = Arc::new(Int32Array::from(vec![
- Some(1),
- Some(1),
- None,
- Some(2),
- Some(2),
- ]));
- ree_write_read_roundtrip(ree, flat);
- }
-
- #[test]
- fn ree_bool() {
- // run_ends [3, 5, 7] → [T,T,T, null,null, F,F]
- let ree: ArrayRef = Arc::new(
- RunArray::try_new(
- &Int32Array::from(vec![3, 5, 7]),
- &BooleanArray::from(vec![Some(true), None, Some(false)]),
- )
- .unwrap(),
- );
- let flat: ArrayRef = Arc::new(BooleanArray::from(vec![
- Some(true),
- Some(true),
- Some(true),
- None,
- None,
- Some(false),
- Some(false),
- ]));
- ree_write_read_roundtrip(ree, flat);
- }
-
- #[test]
- fn ree_fixed_size_binary() {
- let mk = |vals: &[Option<&[u8]>]| -> FixedSizeBinaryArray {
- let mut b = FixedSizeBinaryBuilder::new(2);
- for v in vals {
- match v {
- Some(x) => b.append_value(x).unwrap(),
- None => b.append_null(),
- }
- }
- b.finish()
- };
- // run_ends [2, 4, 6] → [aa,aa, null,null, bb,bb]
- let ree: ArrayRef = Arc::new(
- RunArray::try_new(
- &Int32Array::from(vec![2, 4, 6]),
- &mk(&[Some(b"aa"), None, Some(b"bb")]),
- )
- .unwrap(),
- );
- let flat: ArrayRef = Arc::new(mk(&[
- Some(b"aa"),
- Some(b"aa"),
- None,
- None,
- Some(b"bb"),
- Some(b"bb"),
- ]));
- ree_write_read_roundtrip(ree, flat);
- }
-
- #[test]
- fn ree_single_run() {
- let ree: ArrayRef = Arc::new(["x", "x", "x"].into_iter().collect::<Int32RunArray>());
- let flat: ArrayRef = Arc::new(StringArray::from(vec!["x", "x", "x"]));
- ree_write_read_roundtrip(ree, flat);
- }
-
- #[test]
- fn ree_float32() {
- // run_ends [2, 4, 5] → [1.0, 1.0, null, null, 2.5]
- let ree: ArrayRef = Arc::new(
- RunArray::try_new(
- &Int32Array::from(vec![2, 4, 5]),
- &Float32Array::from(vec![Some(1.0_f32), None, Some(2.5_f32)]),
- )
- .unwrap(),
- );
- let flat: ArrayRef = Arc::new(Float32Array::from(vec![
- Some(1.0_f32),
- Some(1.0_f32),
- None,
- None,
- Some(2.5_f32),
- ]));
- ree_write_read_roundtrip(ree, flat);
- }
-
- #[test]
- fn ree_sliced() {
- // A sliced (non-zero offset) REE array: verify that get_physical_index
- // correctly accounts for the logical offset when expanding.
- // Full array: run_ends [3, 5, 7] → [a,a,a, b,b, c,c]
- // After slice(2, 5) the logical view is [a, b, b, c, c].
- let full: ArrayRef = Arc::new(
- RunArray::try_new(
- &Int32Array::from(vec![3, 5, 7]),
- &StringArray::from(vec!["a", "b", "c"]),
- )
- .unwrap(),
- );
- let sliced = full.slice(2, 5);
- let flat: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "b", "c", "c"]));
- ree_write_read_roundtrip(sliced, flat);
- }
-
#[test]
#[cfg_attr(miri, ignore)] // Takes too long
fn test_number_distinct_values_exact_count() {
diff --git a/parquet/tests/arrow_writer/mod.rs b/parquet/tests/arrow_writer/mod.rs
index c2803c5a13..caf89d3641 100644
--- a/parquet/tests/arrow_writer/mod.rs
+++ b/parquet/tests/arrow_writer/mod.rs
@@ -19,8 +19,6 @@
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;
diff --git a/parquet/tests/arrow_writer/roundtrip.rs b/parquet/tests/arrow_writer/roundtrip.rs
index c9199db609..2e0da18580 100644
--- a/parquet/tests/arrow_writer/roundtrip.rs
+++ b/parquet/tests/arrow_writer/roundtrip.rs
@@ -26,7 +26,8 @@ use std::sync::Arc;
use arrow::datatypes::ToByteSlice;
use arrow_array::builder::{
- FixedSizeBinaryBuilder, ListBuilder, PrimitiveDictionaryBuilder, StringViewBuilder,
+ FixedSizeBinaryBuilder, ListBuilder, PrimitiveDictionaryBuilder, PrimitiveRunBuilder,
+ StringViewBuilder,
};
use arrow_array::cast::AsArray;
use arrow_array::types::{
@@ -35,25 +36,28 @@ use arrow_array::types::{
Time32MillisecondType, Time64MicrosecondType, UInt8Type, UInt16Type, UInt32Type,
};
use arrow_array::{
- Array, ArrayRef, BinaryArray, BinaryViewArray, Date32Array, Date64Array, Decimal32Array,
- Decimal64Array, Decimal128Array, Decimal256Array, DictionaryArray, DurationMicrosecondArray,
- DurationMillisecondArray, DurationNanosecondArray, DurationSecondArray, FixedSizeBinaryArray,
- Float16Array, Float32Array, Float64Array, Int8Array, Int16Array, Int32Array,
- Int32DictionaryArray, Int64Array, LargeBinaryArray, LargeListArray, LargeListViewArray,
- LargeStringArray, ListArray, ListViewArray, NullArray, PrimitiveArray, RecordBatch,
- RecordBatchReader, StringArray, StringViewArray, StructArray, Time32MillisecondArray,
- Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray, TimestampMicrosecondArray,
+ Array, ArrayRef, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, Date64Array,
+ Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, DictionaryArray,
+ DurationMicrosecondArray, DurationMillisecondArray, DurationNanosecondArray,
+ DurationSecondArray, FixedSizeBinaryArray, Float16Array, Float32Array, Float64Array, Int8Array,
+ Int16Array, Int32Array, Int32DictionaryArray, Int32RunArray, Int64Array, IntervalDayTimeArray,
+ IntervalYearMonthArray, LargeBinaryArray, LargeListArray, LargeListViewArray, LargeStringArray,
+ ListArray, ListViewArray, NullArray, PrimitiveArray, RecordBatch, RecordBatchReader, RunArray,
+ 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_buffer::{ArrowNativeType, Buffer, IntervalDayTime, NullBuffer, i256};
use arrow_data::{ArrayData, ArrayDataBuilder};
-use arrow_schema::{DataType as ArrowDataType, Field, Fields, Schema, TimeUnit};
+use arrow_schema::{DataType as ArrowDataType, Field, Fields, Schema, SchemaRef, TimeUnit};
use bytes::Bytes;
use half::f16;
use num_traits::{FromPrimitive, PrimInt, ToPrimitive};
use parquet::arrow::ArrowWriter;
-use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
+use parquet::arrow::arrow_reader::{
+ ArrowReaderOptions, ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder,
+};
use parquet::basic::Type as PhysicalType;
use parquet::errors::Result;
use parquet::file::properties::WriterProperties;
@@ -1231,6 +1235,393 @@ fn arrow_writer_string_dictionary_unsigned_index() {
RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
}
+/// Writes a single-column RecordBatch to an in-memory Parquet buffer.
+fn write_column_to_bytes(array: ArrayRef) -> Bytes {
+ let schema = Arc::new(Schema::new(vec![Field::new(
+ "col",
+ array.data_type().clone(),
+ true,
+ )]));
+ let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
+ let mut buf = Vec::new();
+ let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+ Bytes::from(buf)
+}
+
+/// Reads column 0 from a single-row-group Parquet buffer, projecting it with the given schema.
+/// Passing a flat schema when the buffer was written from a REE array lets callers decode
+/// the physical values without the run-end encoding wrapper.
+fn read_column_with_schema(bytes: Bytes, schema: SchemaRef) -> ArrayRef {
+ let opts = ArrowReaderOptions::new().with_schema(schema);
+ ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, opts)
+ .unwrap()
+ .build()
+ .unwrap()
+ .next()
+ .unwrap()
+ .unwrap()
+ .column(0)
+ .clone()
+}
+
+fn ree_write_read_roundtrip(ree: ArrayRef, flat: ArrayRef) {
+ let flat_schema = Arc::new(Schema::new(vec![Field::new(
+ "col",
+ flat.data_type().clone(),
+ true,
+ )]));
+ let ree_bytes = write_column_to_bytes(ree);
+ let flat_bytes = write_column_to_bytes(flat.clone());
+ assert_eq!(
+ ree_bytes, flat_bytes,
+ "REE and flat bytes should be identical"
+ );
+
+ let decoded_ree = read_column_with_schema(ree_bytes, flat_schema.clone());
+ let decoded_flat = read_column_with_schema(flat_bytes, flat_schema);
+
+ assert_eq!(decoded_ree.as_ref(), flat.as_ref());
+ assert_eq!(decoded_ree.as_ref(), decoded_flat.as_ref());
+}
+
+#[test]
+fn ree_string() {
+ let ree: ArrayRef = Arc::new(
+ [Some("a"), Some("a"), None, Some("b"), Some("b")]
+ .into_iter()
+ .collect::<Int32RunArray>(),
+ );
+ let flat: ArrayRef = Arc::new(StringArray::from(vec![
+ Some("a"),
+ Some("a"),
+ None,
+ Some("b"),
+ Some("b"),
+ ]));
+ ree_write_read_roundtrip(ree, flat);
+}
+
+#[test]
+fn ree_int32() {
+ let mut b = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
+ for v in [Some(1), Some(1), None, Some(2), Some(2)] {
+ b.append_option(v);
+ }
+ let ree: ArrayRef = Arc::new(b.finish());
+ let flat: ArrayRef = Arc::new(Int32Array::from(vec![
+ Some(1),
+ Some(1),
+ None,
+ Some(2),
+ Some(2),
+ ]));
+ ree_write_read_roundtrip(ree, flat);
+}
+
+#[test]
+fn ree_bool() {
+ // run_ends [3, 5, 7] → [T,T,T, null,null, F,F]
+ let ree: ArrayRef = Arc::new(
+ RunArray::try_new(
+ &Int32Array::from(vec![3, 5, 7]),
+ &BooleanArray::from(vec![Some(true), None, Some(false)]),
+ )
+ .unwrap(),
+ );
+ let flat: ArrayRef = Arc::new(BooleanArray::from(vec![
+ Some(true),
+ Some(true),
+ Some(true),
+ None,
+ None,
+ Some(false),
+ Some(false),
+ ]));
+ ree_write_read_roundtrip(ree, flat);
+}
+
+#[test]
+fn ree_fixed_size_binary() {
+ let mk = |vals: &[Option<&[u8]>]| -> FixedSizeBinaryArray {
+ let mut b = FixedSizeBinaryBuilder::new(2);
+ for v in vals {
+ match v {
+ Some(x) => b.append_value(x).unwrap(),
+ None => b.append_null(),
+ }
+ }
+ b.finish()
+ };
+ // run_ends [2, 4, 6] → [aa,aa, null,null, bb,bb]
+ let ree: ArrayRef = Arc::new(
+ RunArray::try_new(
+ &Int32Array::from(vec![2, 4, 6]),
+ &mk(&[Some(b"aa"), None, Some(b"bb")]),
+ )
+ .unwrap(),
+ );
+ let flat: ArrayRef = Arc::new(mk(&[
+ Some(b"aa"),
+ Some(b"aa"),
+ None,
+ None,
+ Some(b"bb"),
+ Some(b"bb"),
+ ]));
+ ree_write_read_roundtrip(ree, flat);
+}
+
+#[test]
+fn ree_single_run() {
+ let ree: ArrayRef = Arc::new(["x", "x", "x"].into_iter().collect::<Int32RunArray>());
+ let flat: ArrayRef = Arc::new(StringArray::from(vec!["x", "x", "x"]));
+ ree_write_read_roundtrip(ree, flat);
+}
+
+#[test]
+fn ree_float32() {
+ // run_ends [2, 4, 5] → [1.0, 1.0, null, null, 2.5]
+ let ree: ArrayRef = Arc::new(
+ RunArray::try_new(
+ &Int32Array::from(vec![2, 4, 5]),
+ &Float32Array::from(vec![Some(1.0_f32), None, Some(2.5_f32)]),
+ )
+ .unwrap(),
+ );
+ let flat: ArrayRef = Arc::new(Float32Array::from(vec![
+ Some(1.0_f32),
+ Some(1.0_f32),
+ None,
+ None,
+ Some(2.5_f32),
+ ]));
+ ree_write_read_roundtrip(ree, flat);
+}
+
+#[test]
+fn ree_sliced() {
+ // A sliced (non-zero offset) REE array: verify that get_physical_index
+ // correctly accounts for the logical offset when expanding.
+ // Full array: run_ends [3, 5, 7] → [a,a,a, b,b, c,c]
+ // After slice(2, 5) the logical view is [a, b, b, c, c].
+ let full: ArrayRef = Arc::new(
+ RunArray::try_new(
+ &Int32Array::from(vec![3, 5, 7]),
+ &StringArray::from(vec!["a", "b", "c"]),
+ )
+ .unwrap(),
+ );
+ let sliced = full.slice(2, 5);
+ let flat: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "b", "c", "c"]));
+ ree_write_read_roundtrip(sliced, flat);
+}
+
+#[test]
+fn arrow_writer() {
+ // define schema
+ let schema = Schema::new(vec![
+ Field::new("a", ArrowDataType::Int32, false),
+ Field::new("b", ArrowDataType::Int32, true),
+ ]);
+
+ // create some data
+ let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
+ let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
+
+ // build a record batch
+ let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
+
+ roundtrip(batch, Some(SMALL_SIZE / 2));
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_non_null() {
+ let schema = Schema::new(vec![Field::new("a", ArrowDataType::Int32, false)]);
+ let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
+
+ RoundTripTest::new(Arc::new(a))
+ .with_schema(Arc::new(schema))
+ .run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_binary() {
+ let raw_string_values = vec!["foo", "bar", "baz", "quux"];
+ let raw_binary_values = [
+ b"foo".to_vec(),
+ b"bar".to_vec(),
+ b"baz".to_vec(),
+ b"quux".to_vec(),
+ ];
+ let raw_binary_value_refs = raw_binary_values
+ .iter()
+ .map(|x| x.as_slice())
+ .collect::<Vec<_>>();
+
+ let string_values = StringArray::from(raw_string_values.clone());
+ let binary_values = BinaryArray::from(raw_binary_value_refs);
+ assert_eq!(string_values.null_count(), 0);
+ assert_eq!(binary_values.null_count(), 0);
+
+ RoundTripTest::new(Arc::new(string_values)).run();
+ RoundTripTest::new(Arc::new(binary_values)).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_binary_view() {
+ let raw_string_values = vec!["foo", "bar", "large payload over 12 bytes", "lulu"];
+ let raw_binary_values = vec![
+ b"foo".to_vec(),
+ b"bar".to_vec(),
+ b"large payload over 12 bytes".to_vec(),
+ b"lulu".to_vec(),
+ ];
+ let nullable_string_values = vec![Some("foo"), None, Some("large payload over 12 bytes"), None];
+
+ 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);
+
+ 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]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn arrow_writer_binary_view_long_value() {
+ // 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.
+ let long = "a".repeat(128);
+ let raw_string_values = vec!["foo", long.as_str(), "bar"];
+ let raw_binary_values = vec![b"foo".to_vec(), long.as_bytes().to_vec(), b"bar".to_vec()];
+
+ let string_view_values: ArrayRef = Arc::new(StringViewArray::from(raw_string_values));
+ let binary_view_values: ArrayRef =
+ Arc::new(BinaryViewArray::from_iter_values(raw_binary_values));
+
+ RoundTripTest::new(Arc::clone(&string_view_values))
+ .with_nullable(false)
+ .run();
+ RoundTripTest::new(Arc::clone(&binary_view_values))
+ .with_nullable(false)
+ .run();
+}
+
+fn get_decimal_batch(precision: u8, scale: i8) -> RecordBatch {
+ let decimal_field = Field::new("a", ArrowDataType::Decimal128(precision, scale), false);
+ let schema = Schema::new(vec![decimal_field]);
+
+ let decimal_values = vec![10_000, 50_000, 0, -100]
+ .into_iter()
+ .map(Some)
+ .collect::<Decimal128Array>()
+ .with_precision_and_scale(precision, scale)
+ .unwrap();
+
+ RecordBatch::try_new(Arc::new(schema), vec![Arc::new(decimal_values)]).unwrap()
+}
+
+#[test]
+fn arrow_writer_decimal() {
+ // int32 to store the decimal value
+ let batch_int32_decimal = get_decimal_batch(5, 2);
+ roundtrip(batch_int32_decimal, Some(SMALL_SIZE / 2));
+ // int64 to store the decimal value
+ let batch_int64_decimal = get_decimal_batch(12, 2);
+ roundtrip(batch_int64_decimal, Some(SMALL_SIZE / 2));
+ // fixed_length_byte_array to store the decimal value
+ let batch_fixed_len_byte_array_decimal = get_decimal_batch(30, 2);
+ roundtrip(batch_fixed_len_byte_array_decimal, Some(SMALL_SIZE / 2));
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // inline assembly is not supported
+fn arrow_writer_float_nans() {
+ const MEDIUM_SIZE: usize = 63;
+
+ let f16_field = Field::new("a", ArrowDataType::Float16, false);
+ let f32_field = Field::new("b", ArrowDataType::Float32, false);
+ let f64_field = Field::new("c", ArrowDataType::Float64, false);
+ let schema = Schema::new(vec![f16_field, f32_field, f64_field]);
+
+ let f16_values = (0..MEDIUM_SIZE)
+ .map(|i| {
+ Some(if i % 2 == 0 {
+ f16::NAN
+ } else {
+ f16::from_f32(i as f32)
+ })
+ })
+ .collect::<Float16Array>();
+
+ let f32_values = (0..MEDIUM_SIZE)
+ .map(|i| Some(if i % 2 == 0 { f32::NAN } else { i as f32 }))
+ .collect::<Float32Array>();
+
+ let f64_values = (0..MEDIUM_SIZE)
+ .map(|i| Some(if i % 2 == 0 { f64::NAN } else { i as f64 }))
+ .collect::<Float64Array>();
+
+ let batch = RecordBatch::try_new(
+ Arc::new(schema),
+ vec![
+ Arc::new(f16_values),
+ Arc::new(f32_values),
+ Arc::new(f64_values),
+ ],
+ )
+ .unwrap();
+
+ roundtrip(batch, None);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn all_null_primitive_single_column() {
+ let values = Arc::new(Int32Array::from(vec![None; SMALL_SIZE]));
+ RoundTripTest::new(values).run();
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn null_single_column() {
+ let values = Arc::new(NullArray::new(SMALL_SIZE));
+ RoundTripTest::new(values).run();
+ // null arrays are always nullable, a test with non-nullable nulls fails
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn bool_single_column() {
+ required_and_optional::<BooleanArray, _>(
+ [true, false].iter().cycle().copied().take(SMALL_SIZE),
+ );
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn interval_year_month_single_column() {
+ required_and_optional::<IntervalYearMonthArray, _>(0..SMALL_SIZE as i32);
+}
+
+#[test]
+#[cfg_attr(miri, ignore)] // Takes too long
+fn interval_day_time_single_column() {
+ required_and_optional::<IntervalDayTimeArray, _>(vec![
+ IntervalDayTime::new(0, 1),
+ IntervalDayTime::new(0, 3),
+ IntervalDayTime::new(3, -2),
+ IntervalDayTime::new(-200, 4),
+ ]);
+}
+
#[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
index 5d541e227a..9851d3194e 100644
--- a/parquet/tests/arrow_writer/roundtrip_helpers.rs
+++ b/parquet/tests/arrow_writer/roundtrip_helpers.rs
@@ -35,6 +35,9 @@ 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
+// expect(dead_code) cannot apply to both integration and unit tests with different usage.
+#[expect(clippy::allow_attributes)]
+#[allow(dead_code)]
pub(super) fn roundtrip(
expected_batch: RecordBatch,
max_row_group_size: Option<usize>,
@@ -127,31 +130,43 @@ impl RoundTripTest {
}
}
- /// 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 the schema
+ // expect(dead_code) cannot apply to both integration and unit tests with different usage.
+ #[expect(clippy::allow_attributes)]
+ #[allow(dead_code)]
+ pub(super) fn with_schema(mut self, schema: SchemaRef) -> Self {
+ self.schema = Some(schema);
+ self
+ }
+
/// Set bloom filter
+ // expect(dead_code) cannot apply to both integration and unit tests with different usage.
+ #[expect(clippy::allow_attributes)]
+ #[allow(dead_code)]
pub(super) fn with_bloom_filter(mut self, bloom_filter: bool) -> Self {
self.bloom_filter = bloom_filter;
self
}
/// Set bloom filter max ndv
+ // expect(dead_code) cannot apply to both integration and unit tests with different usage.
+ #[expect(clippy::allow_attributes)]
+ #[allow(dead_code)]
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
+ // expect(dead_code) cannot apply to both integration and unit tests with different usage.
+ #[expect(clippy::allow_attributes)]
+ #[allow(dead_code)]
pub(super) fn with_bloom_filter_position(
mut self,
bloom_filter_position: BloomFilterPosition,