diff --git a/arrow-string/src/concat_elements.rs b/arrow-string/src/concat_elements.rs index 8493dfafe0aa..8def0d5abb10 100644 --- a/arrow-string/src/concat_elements.rs +++ b/arrow-string/src/concat_elements.rs @@ -402,23 +402,26 @@ pub fn concat_elements_string_view_array( /// Returns the elementwise concatenation of [`Array`]s. /// +/// The output array will have the same type as the input arrays (which must have the same type). +/// +/// Concatenation of `FixedSizeBinaryArray` instances with different sizes is supported. The output +/// type is `FixedSizeBinaryArray` with the sum of the sizes of the two input arrays as size. +/// /// # Errors /// /// This function errors if the arrays are of different types. pub fn concat_elements_dyn(left: &dyn Array, right: &dyn Array) -> Result { - if left.data_type() != right.data_type() { - return Err(ArrowError::ComputeError(format!( - "Cannot concat arrays of different types: {} != {}", - left.data_type(), - right.data_type() - ))); - } match (left.data_type(), right.data_type()) { (DataType::Utf8, DataType::Utf8) => { let left = left.as_any().downcast_ref::().unwrap(); let right = right.as_any().downcast_ref::().unwrap(); Ok(Arc::new(concat_elements_utf8(left, right)?)) } + (DataType::Utf8View, DataType::Utf8View) => { + let left = left.as_any().downcast_ref::().unwrap(); + let right = right.as_any().downcast_ref::().unwrap(); + Ok(Arc::new(concat_elements_string_view_array(left, right)?)) + } (DataType::LargeUtf8, DataType::LargeUtf8) => { let left = left.as_any().downcast_ref::().unwrap(); let right = right.as_any().downcast_ref::().unwrap(); @@ -429,20 +432,15 @@ pub fn concat_elements_dyn(left: &dyn Array, right: &dyn Array) -> Result().unwrap(); Ok(Arc::new(concat_element_binary(left, right)?)) } - (DataType::LargeBinary, DataType::LargeBinary) => { - let left = left.as_any().downcast_ref::().unwrap(); - let right = right.as_any().downcast_ref::().unwrap(); - Ok(Arc::new(concat_element_binary(left, right)?)) - } (DataType::BinaryView, DataType::BinaryView) => { let left = left.as_any().downcast_ref::().unwrap(); let right = right.as_any().downcast_ref::().unwrap(); Ok(Arc::new(concat_elements_binary_view_array(left, right)?)) } - (DataType::Utf8View, DataType::Utf8View) => { - let left = left.as_any().downcast_ref::().unwrap(); - let right = right.as_any().downcast_ref::().unwrap(); - Ok(Arc::new(concat_elements_string_view_array(left, right)?)) + (DataType::LargeBinary, DataType::LargeBinary) => { + let left = left.as_any().downcast_ref::().unwrap(); + let right = right.as_any().downcast_ref::().unwrap(); + Ok(Arc::new(concat_element_binary(left, right)?)) } (DataType::FixedSizeBinary(_), DataType::FixedSizeBinary(_)) => { let left = left @@ -456,16 +454,26 @@ pub fn concat_elements_dyn(left: &dyn Array, right: &dyn Array) -> Result Err(ArrowError::NotYetImplemented(format!( - "concat not supported for {}", - left.data_type() - ))), + (l, r) => { + if l != r { + Err(ArrowError::ComputeError(format!( + "Cannot concat arrays of different types: {} != {}", + l, r + ))) + } else { + Err(ArrowError::NotYetImplemented(format!( + "concat not supported for {}", + left.data_type() + ))) + } + } } } #[cfg(test)] mod tests { use super::*; + use arrow_array::cast::AsArray; use arrow_buffer::Buffer; #[test] @@ -477,13 +485,13 @@ mod tests { .into_iter() .collect::(); - let output = concat_elements_utf8(&left, &right).unwrap(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = [None, Some("baryyy"), None] .into_iter() .collect::(); - assert_eq!(output, expected); + assert_eq!(output.as_string(), &expected); } #[test] @@ -495,13 +503,13 @@ mod tests { .into_iter() .collect::(); - let output = concat_elements_utf8(&left, &right).unwrap(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = [Some("foobaz"), Some(""), Some("bar")] .into_iter() .collect::(); - assert_eq!(output, expected); + assert_eq!(output.as_string(), &expected); } #[test] @@ -509,11 +517,11 @@ mod tests { let left = StringArray::from(vec!["foo", "bar"]); let right = StringArray::from(vec!["bar", "baz"]); - let output = concat_elements_utf8(&left, &right).unwrap(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = StringArray::from(vec!["foobar", "barbaz"]); - assert_eq!(output, expected); + assert_eq!(output.as_string(), &expected); } #[test] @@ -521,7 +529,7 @@ mod tests { let left = StringArray::from(vec!["foo", "bar"]); let right = StringArray::from(vec!["baz"]); - let output = concat_elements_utf8(&left, &right); + let output = concat_elements_dyn(&left, &right); assert_eq!( output.unwrap_err().to_string(), @@ -536,7 +544,7 @@ mod tests { let left_slice = left.slice(0, 3); let right_slice = right.slice(1, 3); - let output = concat_elements_utf8( + let output = concat_elements_dyn( left_slice .as_any() .downcast_ref::>() @@ -552,12 +560,12 @@ mod tests { .into_iter() .collect::(); - assert_eq!(output, expected); + assert_eq!(output.as_string(), &expected); let left_slice = left.slice(2, 2); let right_slice = right.slice(1, 2); - let output = concat_elements_utf8( + let output = concat_elements_dyn( left_slice .as_any() .downcast_ref::>() @@ -571,7 +579,7 @@ mod tests { let expected = [None, Some("bazfar")].into_iter().collect::(); - assert_eq!(output, expected); + assert_eq!(output.as_string(), &expected); } #[test] @@ -617,11 +625,26 @@ mod tests { let right = FixedSizeBinaryArray::try_from(vec![None, Some(b"yyy" as &[u8]), Some(b"zzz")]) .unwrap(); - let output = concat_elements_fixed_size_binary(&left, &right).unwrap(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = FixedSizeBinaryArray::try_from(vec![None, Some(b"baryyy" as &[u8]), None]).unwrap(); - assert_eq!(output, expected); + assert_eq!(output.as_fixed_size_binary(), &expected); + } + + #[test] + fn test_mixed_fixed_size_binary_concat() { + let left = + FixedSizeBinaryArray::try_from(vec![Some(b"foobar" as &[u8]), Some(b"barbaz"), None]) + .unwrap(); + let right = FixedSizeBinaryArray::try_from(vec![None, Some(b"yyy" as &[u8]), Some(b"zzz")]) + .unwrap(); + + let output = concat_elements_dyn(&left, &right).unwrap(); + + let expected = + FixedSizeBinaryArray::try_from(vec![None, Some(b"barbazyyy" as &[u8]), None]).unwrap(); + assert_eq!(output.as_fixed_size_binary(), &expected); } #[test] @@ -629,10 +652,10 @@ mod tests { let left = FixedSizeBinaryArray::try_from(vec![b"ab" as &[u8], b"cd"]).unwrap(); let right = FixedSizeBinaryArray::try_from(vec![b"12" as &[u8], b"34"]).unwrap(); - let output = concat_elements_fixed_size_binary(&left, &right).unwrap(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = FixedSizeBinaryArray::try_from(vec![b"ab12" as &[u8], b"cd34"]).unwrap(); - assert_eq!(output, expected); + assert_eq!(output.as_fixed_size_binary(), &expected); } #[test] @@ -640,7 +663,7 @@ mod tests { let left = FixedSizeBinaryArray::try_from(vec![b"ab" as &[u8], b"cd"]).unwrap(); let right = FixedSizeBinaryArray::try_from(vec![b"12" as &[u8]]).unwrap(); - let output = concat_elements_fixed_size_binary(&left, &right); + let output = concat_elements_dyn(&left, &right); assert_eq!( output.unwrap_err().to_string(), "Compute error: Arrays must have the same length: 2 != 1".to_string() @@ -652,10 +675,10 @@ mod tests { let left = FixedSizeBinaryArray::new(0, Buffer::from(&[]), None); let right = FixedSizeBinaryArray::new(0, Buffer::from(&[]), None); - let output = concat_elements_fixed_size_binary(&left, &right).unwrap(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = FixedSizeBinaryArray::new(0, Buffer::from(&[]), None); - assert_eq!(output, expected); + assert_eq!(output.as_fixed_size_binary(), &expected); } #[test] @@ -716,7 +739,7 @@ mod tests { Some(long), ]); - let output = concat_elements_string_view_array(&left, &right).unwrap(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = StringViewArray::from_iter(vec![ None, @@ -727,7 +750,7 @@ mod tests { Some("ThisStringIsLongerThan12Bytesbar"), Some("ThisStringIsLongerThan12BytesThisStringIsLongerThan12Bytes"), ]); - assert_eq!(output, expected); + assert_eq!(output.as_string_view(), &expected); let left = StringViewArray::from_iter(vec![ Some("a"), @@ -746,7 +769,7 @@ mod tests { Some(long), ]); - let output = concat_elements_string_view_array(&left, &right).unwrap(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = StringViewArray::from_iter(vec![ Some("ac"), @@ -756,7 +779,7 @@ mod tests { Some("ThisStringIsLongerThan12Bytesd"), Some("ThisStringIsLongerThan12BytesThisStringIsLongerThan12Bytes"), ]); - assert_eq!(output, expected); + assert_eq!(output.as_string_view(), &expected); } #[test] @@ -774,7 +797,7 @@ mod tests { Some(b""), ]); - let output = concat_elements_binary_view_array(&left, &right).unwrap(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = BinaryViewArray::from_iter(vec![ Some(b"foobar" as &[u8]), @@ -782,7 +805,7 @@ mod tests { Some(b""), Some(b"baz"), ]); - assert_eq!(output, expected); + assert_eq!(output.as_binary_view(), &expected); } #[test] @@ -790,7 +813,7 @@ mod tests { let left = BinaryViewArray::from_iter(vec![Some(b"foo" as &[u8]), Some(b"bar")]); let right = BinaryViewArray::from_iter(vec![Some(b"baz" as &[u8])]); - let output = concat_elements_binary_view_array(&left, &right); + let output = concat_elements_dyn(&left, &right); assert_eq!( output.unwrap_err().to_string(), "Compute error: Arrays must have the same length: 2 != 1".to_string() @@ -802,9 +825,9 @@ mod tests { let left = BinaryViewArray::from_iter(vec![] as Vec>); let right = BinaryViewArray::from_iter(vec![] as Vec>); - let output = concat_elements_binary_view_array(&left, &right).unwrap(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = BinaryViewArray::from_iter(vec![] as Vec>); - assert_eq!(output, expected); + assert_eq!(output.as_binary_view(), &expected); } #[test] @@ -813,43 +836,31 @@ mod tests { let left = StringArray::from(vec![Some("foo"), Some("bar"), None]); let right = StringArray::from(vec![None, Some("yyy"), Some("zzz")]); - let output: StringArray = concat_elements_dyn(&left, &right) - .unwrap() - .into_data() - .into(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = StringArray::from(vec![None, Some("baryyy"), None]); - assert_eq!(output, expected); + assert_eq!(output.as_string(), &expected); // test for LargeStringArray let left = LargeStringArray::from(vec![Some("foo"), Some("bar"), None]); let right = LargeStringArray::from(vec![None, Some("yyy"), Some("zzz")]); - let output: LargeStringArray = concat_elements_dyn(&left, &right) - .unwrap() - .into_data() - .into(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = LargeStringArray::from(vec![None, Some("baryyy"), None]); - assert_eq!(output, expected); + assert_eq!(output.as_string(), &expected); // test for BinaryArray let left = BinaryArray::from_opt_vec(vec![Some(b"foo"), Some(b"bar"), None]); let right = BinaryArray::from_opt_vec(vec![None, Some(b"yyy"), Some(b"zzz")]); - let output: BinaryArray = concat_elements_dyn(&left, &right) - .unwrap() - .into_data() - .into(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = BinaryArray::from_opt_vec(vec![None, Some(b"baryyy"), None]); - assert_eq!(output, expected); + assert_eq!(output.as_binary(), &expected); // test for LargeBinaryArray let left = LargeBinaryArray::from_opt_vec(vec![Some(b"foo"), Some(b"bar"), None]); let right = LargeBinaryArray::from_opt_vec(vec![None, Some(b"yyy"), Some(b"zzz")]); - let output: LargeBinaryArray = concat_elements_dyn(&left, &right) - .unwrap() - .into_data() - .into(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = LargeBinaryArray::from_opt_vec(vec![None, Some(b"baryyy"), None]); - assert_eq!(output, expected); + assert_eq!(output.as_binary(), &expected); // test for BinaryViewArray let long = b"ThisStringIsLongerThan12Bytes" as &[u8]; @@ -871,10 +882,7 @@ mod tests { Some(b"bar"), Some(long), ]); - let output: BinaryViewArray = concat_elements_dyn(&left, &right) - .unwrap() - .into_data() - .into(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = BinaryViewArray::from_iter(vec![ None, Some(b"baryyy" as &[u8]), @@ -884,7 +892,7 @@ mod tests { Some(b"ThisStringIsLongerThan12Bytesbar"), Some(b"ThisStringIsLongerThan12BytesThisStringIsLongerThan12Bytes"), ]); - assert_eq!(output, expected); + assert_eq!(output.as_binary_view(), &expected); // test for StringViewArray let long = "ThisStringIsLongerThan12Bytes"; @@ -906,10 +914,7 @@ mod tests { Some("bar"), Some(long), ]); - let output: StringViewArray = concat_elements_dyn(&left, &right) - .unwrap() - .into_data() - .into(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = StringViewArray::from_iter(vec![ None, Some("baryyy"), @@ -919,20 +924,17 @@ mod tests { Some("ThisStringIsLongerThan12Bytesbar"), Some("ThisStringIsLongerThan12BytesThisStringIsLongerThan12Bytes"), ]); - assert_eq!(output, expected); + assert_eq!(output.as_string_view(), &expected); // test for FixedSizeBinaryArray let left = FixedSizeBinaryArray::try_from(vec![Some(b"foo" as &[u8]), Some(b"bar"), None]) .unwrap(); let right = FixedSizeBinaryArray::try_from(vec![None, Some(b"yyy" as &[u8]), Some(b"zzz")]) .unwrap(); - let output: FixedSizeBinaryArray = concat_elements_dyn(&left, &right) - .unwrap() - .into_data() - .into(); + let output = concat_elements_dyn(&left, &right).unwrap(); let expected = FixedSizeBinaryArray::try_from(vec![None, Some(b"baryyy" as &[u8]), None]).unwrap(); - assert_eq!(output, expected); + assert_eq!(output.as_fixed_size_binary(), &expected); } #[test]