diff --git a/rust/arrow/src/array/transform/mod.rs b/rust/arrow/src/array/transform/mod.rs index 689118d20f86..23d9dfe6c400 100644 --- a/rust/arrow/src/array/transform/mod.rs +++ b/rust/arrow/src/array/transform/mod.rs @@ -714,6 +714,46 @@ mod tests { assert_eq!(array, expected) } + #[test] + fn test_struct_offset() { + let strings: ArrayRef = Arc::new(StringArray::from(vec![ + Some("joe"), + None, + None, + Some("mark"), + Some("doe"), + ])); + let ints: ArrayRef = Arc::new(Int32Array::from(vec![ + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + ])); + + let array = + StructArray::try_from(vec![("f1", strings.clone()), ("f2", ints.clone())]) + .unwrap() + .slice(1, 3) + .data(); + let arrays = vec![array.as_ref()]; + let mut mutable = MutableArrayData::new(arrays, false, 0); + + mutable.extend(0, 1, 3); + let data = mutable.freeze(); + let array = StructArray::from(Arc::new(data)); + + let expected_strings: ArrayRef = + Arc::new(StringArray::from(vec![None, Some("mark")])); + let expected = StructArray::try_from(vec![ + ("f1", expected_strings), + ("f2", ints.slice(2, 2)), + ]) + .unwrap(); + + assert_eq!(array, expected); + } + #[test] fn test_struct_nulls() { let strings: ArrayRef = Arc::new(StringArray::from(vec![ diff --git a/rust/arrow/src/array/transform/structure.rs b/rust/arrow/src/array/transform/structure.rs index 5c41d76a7f1c..c019f5ac6a93 100644 --- a/rust/arrow/src/array/transform/structure.rs +++ b/rust/arrow/src/array/transform/structure.rs @@ -26,10 +26,13 @@ pub(super) fn build_extend(array: &ArrayData) -> Extend { index: usize, start: usize, len: usize| { - mutable - .child_data - .iter_mut() - .for_each(|child| child.extend(index, start, start + len)) + mutable.child_data.iter_mut().for_each(|child| { + child.extend( + index, + array.offset() + start, + array.offset() + start + len, + ) + }) }, ) } else { @@ -38,7 +41,7 @@ pub(super) fn build_extend(array: &ArrayData) -> Extend { index: usize, start: usize, len: usize| { - (start..start + len).for_each(|i| { + (array.offset() + start..array.offset() + start + len).for_each(|i| { if array.is_valid(i) { mutable .child_data diff --git a/rust/arrow/src/compute/kernels/concat.rs b/rust/arrow/src/compute/kernels/concat.rs index 08240424e822..a51831c744ef 100644 --- a/rust/arrow/src/compute/kernels/concat.rs +++ b/rust/arrow/src/compute/kernels/concat.rs @@ -157,6 +157,40 @@ mod tests { Ok(()) } + #[test] + fn test_concat_primitive_array_slices() -> Result<()> { + let input_1 = PrimitiveArray::::from(vec![ + Some(-1), + Some(-1), + Some(2), + None, + None, + ]) + .slice(1, 3); + + let input_2 = PrimitiveArray::::from(vec![ + Some(101), + Some(102), + Some(103), + None, + ]) + .slice(1, 3); + let arr = concat(&[input_1.as_ref(), input_2.as_ref()])?; + + let expected_output = Arc::new(PrimitiveArray::::from(vec![ + Some(-1), + Some(2), + None, + Some(102), + Some(103), + None, + ])) as ArrayRef; + + assert_eq!(&arr, &expected_output); + + Ok(()) + } + #[test] fn test_concat_boolean_primitive_arrays() -> Result<()> { let arr = concat(&[ @@ -254,4 +288,106 @@ mod tests { Ok(()) } + + #[test] + fn test_concat_struct_arrays() -> Result<()> { + let field = Field::new("field", DataType::Int64, true); + let input_primitive_1: ArrayRef = + Arc::new(PrimitiveArray::::from(vec![ + Some(-1), + Some(-1), + Some(2), + None, + None, + ])); + let input_struct_1 = StructArray::from(vec![(field.clone(), input_primitive_1)]); + + let input_primitive_2: ArrayRef = + Arc::new(PrimitiveArray::::from(vec![ + Some(101), + Some(102), + Some(103), + None, + ])); + let input_struct_2 = StructArray::from(vec![(field.clone(), input_primitive_2)]); + + let input_primitive_3: ArrayRef = + Arc::new(PrimitiveArray::::from(vec![ + Some(256), + Some(512), + Some(1024), + ])); + let input_struct_3 = StructArray::from(vec![(field, input_primitive_3)]); + + let arr = concat(&[&input_struct_1, &input_struct_2, &input_struct_3])?; + + let expected_primitive_output = Arc::new(PrimitiveArray::::from(vec![ + Some(-1), + Some(-1), + Some(2), + None, + None, + Some(101), + Some(102), + Some(103), + None, + Some(256), + Some(512), + Some(1024), + ])) as ArrayRef; + + let actual_primitive = arr + .as_any() + .downcast_ref::() + .unwrap() + .column(0); + assert_eq!(actual_primitive, &expected_primitive_output); + + Ok(()) + } + + #[test] + fn test_concat_struct_array_slices() -> Result<()> { + let field = Field::new("field", DataType::Int64, true); + let input_primitive_1: ArrayRef = + Arc::new(PrimitiveArray::::from(vec![ + Some(-1), + Some(-1), + Some(2), + None, + None, + ])); + let input_struct_1 = StructArray::from(vec![(field.clone(), input_primitive_1)]); + + let input_primitive_2: ArrayRef = + Arc::new(PrimitiveArray::::from(vec![ + Some(101), + Some(102), + Some(103), + None, + ])); + let input_struct_2 = StructArray::from(vec![(field, input_primitive_2)]); + + let arr = concat(&[ + input_struct_1.slice(1, 3).as_ref(), + input_struct_2.slice(1, 2).as_ref(), + ])?; + + let expected_primitive_output = Arc::new(PrimitiveArray::::from(vec![ + Some(-1), + Some(2), + None, + Some(102), + Some(103), + ])) as ArrayRef; + + let actual_primitive = arr + .as_any() + .downcast_ref::() + .unwrap() + .column(0); + assert_eq!(actual_primitive, &expected_primitive_output); + + Ok(()) + } }