From e69d41c9540e23274304c43fd81d6b83d3e9e491 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Tue, 9 Jun 2026 15:12:30 -0400 Subject: [PATCH] Add interleave tests for List and List These exercise interleave over list arrays whose primitive child carries logical type parameters (Decimal precision/scale, Timestamp timezone), asserting the interleaved child preserves the parameterized DataType and correctly carries child-element nulls. Both i32 and i64 offsets covered. Co-Authored-By: Claude Opus 4.8 (1M context) --- arrow-select/src/interleave.rs | 105 ++++++++++++++++++++++++++++++++- 1 file changed, 102 insertions(+), 3 deletions(-) diff --git a/arrow-select/src/interleave.rs b/arrow-select/src/interleave.rs index 13d28747dfcf..542190ab77e1 100644 --- a/arrow-select/src/interleave.rs +++ b/arrow-select/src/interleave.rs @@ -761,10 +761,12 @@ pub fn interleave_record_batch( mod tests { use super::*; use arrow_array::Int32RunArray; - use arrow_array::builder::{GenericListBuilder, Int32Builder, PrimitiveRunBuilder}; - use arrow_array::types::Int8Type; + use arrow_array::builder::{ + GenericListBuilder, Int32Builder, PrimitiveBuilder, PrimitiveRunBuilder, + }; + use arrow_array::types::{Decimal128Type, Int8Type, TimestampMicrosecondType}; use arrow_buffer::ScalarBuffer; - use arrow_schema::Field; + use arrow_schema::{Field, TimeUnit}; #[test] fn test_primitive() { @@ -951,6 +953,103 @@ mod tests { test_interleave_lists::(); } + /// One list slot in a `List` fixture: `None` is a null slot, + /// `Some(items)` is a list whose items may individually be null. + type ListRow = Option::Native>>>; + + /// Build a `List` from row fixtures. The primitive child carries + /// `data_type` (e.g. its Decimal scale or timezone). + fn list_of_primitive( + data_type: &DataType, + rows: &[ListRow], + ) -> GenericListArray { + let mut builder = GenericListBuilder::::new( + PrimitiveBuilder::::new().with_data_type(data_type.clone()), + ); + for row in rows { + match row { + Some(items) => { + items + .iter() + .for_each(|v| builder.values().append_option(*v)); + builder.append(true); + } + None => builder.append(false), + } + } + builder.finish() + } + + /// Interleave list fixtures and assert both the result and that the + /// interleaved primitive child preserves the parameterized `data_type`. + fn check_interleave_list_primitive( + data_type: &DataType, + inputs: &[&[ListRow]], + indices: &[(usize, usize)], + expected: &[ListRow], + ) { + let arrays: Vec<_> = inputs + .iter() + .map(|rows| list_of_primitive::(data_type, rows)) + .collect(); + let refs: Vec<&dyn Array> = arrays.iter().map(|a| a as &dyn Array).collect(); + + let values = interleave(&refs, indices).unwrap(); + let v = values + .as_any() + .downcast_ref::>() + .unwrap(); + + assert_eq!(v, &list_of_primitive::(data_type, expected)); + // The child's logical type (Decimal precision/scale, Timestamp timezone) + // must be preserved, not reset to the primitive's default. + assert_eq!(v.values().data_type(), data_type); + } + + fn test_interleave_lists_decimal() { + // List, exercising child-element nulls and null slots. + check_interleave_list_primitive::( + &DataType::Decimal128(20, 3), + &[ + &[ + Some(vec![Some(1), Some(2)]), + None, + Some(vec![Some(3), None]), + ], // a + &[Some(vec![Some(4)]), Some(vec![Some(5), Some(6)])], // b + ], + &[(0, 2), (0, 1), (1, 0), (1, 1)], + &[ + Some(vec![Some(3), None]), + None, + Some(vec![Some(4)]), + Some(vec![Some(5), Some(6)]), + ], + ); + } + + #[test] + fn test_lists_decimal() { + test_interleave_lists_decimal::(); + test_interleave_lists_decimal::(); + } + + fn test_interleave_lists_timestamp_tz() { + // List, checking the timezone survives. + check_interleave_list_primitive::( + &DataType::Timestamp(TimeUnit::Microsecond, Some("+08:00".into())), + &[&[Some(vec![Some(1), Some(2)]), Some(vec![Some(3)])]], + &[(0, 1), (0, 0)], + &[Some(vec![Some(3)]), Some(vec![Some(1), Some(2)])], + ); + } + + #[test] + fn test_lists_timestamp_tz() { + test_interleave_lists_timestamp_tz::(); + test_interleave_lists_timestamp_tz::(); + } + fn test_interleave_list_views() { // [[1, 2], null, [3]] let mut a = GenericListBuilder::::new(Int32Builder::new());