ARROW-11394: [Rust] Tests for Slice & Concat

I took a stab at adding some tests for concat (and the underlying MutableArrayData) that cover the use case of the array having an offset. I had some problems with the assertions on the string array which makes me think this isn't fully fixed, but I wanted to put it up and see if anyone had suggestions for how to address this. Closes #9339 from bjchambers/ARROW-11394-slice-concat Authored-by: Ben Chambers <bchambers@apache.org> Signed-off-by: Jorge C. Leitao <jorgecarleitao@gmail.com>
apache · Jan 31, 2021 · 9f96561 · 9f96561
1 parent f055d5e
commit 9f96561
Show file tree

Hide file tree

Showing 3 changed files with 184 additions and 5 deletions.
diff --git 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
@@ -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
@@ -157,6 +157,40 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_concat_primitive_array_slices() -> Result<()> {
+        let input_1 = PrimitiveArray::<Int64Type>::from(vec![
+            Some(-1),
+            Some(-1),
+            Some(2),
+            None,
+            None,
+        ])
+        .slice(1, 3);
+
+        let input_2 = PrimitiveArray::<Int64Type>::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::<Int64Type>::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::<Int64Type>::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::<Int64Type>::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::<Int64Type>::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::<Int64Type>::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::<StructArray>()
+            .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::<Int64Type>::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::<Int64Type>::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::<Int64Type>::from(vec![
+            Some(-1),
+            Some(2),
+            None,
+            Some(102),
+            Some(103),
+        ])) as ArrayRef;
+
+        let actual_primitive = arr
+            .as_any()
+            .downcast_ref::<StructArray>()
+            .unwrap()
+            .column(0);
+        assert_eq!(actual_primitive, &expected_primitive_output);
+
+        Ok(())
+    }
 }