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 <[email protected]> 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,
