You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by jo...@apache.org on 2021/01/31 14:37:55 UTC
[arrow] branch master updated: ARROW-11394: [Rust] Tests for Slice
& Concat
This is an automated email from the ASF dual-hosted git repository.
jorgecarleitao pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push:
new 77ae93d ARROW-11394: [Rust] Tests for Slice & Concat
77ae93d is described below
commit 77ae93d6ecaac8fb5f4a18ca5287b7456cd88784
Author: Ben Chambers <bc...@apache.org>
AuthorDate: Sun Jan 31 15:36:51 2021 +0100
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 <bc...@apache.org>
Signed-off-by: Jorge C. Leitao <jo...@gmail.com>
---
rust/arrow/src/array/transform/mod.rs | 40 ++++++++
rust/arrow/src/array/transform/structure.rs | 13 ++-
rust/arrow/src/compute/kernels/concat.rs | 136 ++++++++++++++++++++++++++++
3 files changed, 184 insertions(+), 5 deletions(-)
diff --git a/rust/arrow/src/array/transform/mod.rs b/rust/arrow/src/array/transform/mod.rs
index 689118d..23d9dfe 100644
--- a/rust/arrow/src/array/transform/mod.rs
+++ b/rust/arrow/src/array/transform/mod.rs
@@ -715,6 +715,46 @@ mod tests {
}
#[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![
Some("joe"),
diff --git a/rust/arrow/src/array/transform/structure.rs b/rust/arrow/src/array/transform/structure.rs
index 5c41d76..c019f5a 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 0824042..a51831c 100644
--- a/rust/arrow/src/compute/kernels/concat.rs
+++ b/rust/arrow/src/compute/kernels/concat.rs
@@ -158,6 +158,40 @@ mod tests {
}
#[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(&[
&BooleanArray::from(vec![
@@ -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(())
+ }
}