diff --git a/arrow-string/src/concat_elements.rs b/arrow-string/src/concat_elements.rs index 41be8a81cb12..cd4676d28752 100644 --- a/arrow-string/src/concat_elements.rs +++ b/arrow-string/src/concat_elements.rs @@ -18,10 +18,12 @@ //! Provides utility functions for concatenation of elements in arrays. use std::sync::Arc; -use arrow_array::builder::BufferBuilder; +use arrow_array::builder::{ + BinaryViewBuilder, BufferBuilder, FixedSizeBinaryBuilder, StringViewBuilder, +}; use arrow_array::types::ByteArrayType; use arrow_array::*; -use arrow_buffer::{ArrowNativeType, NullBuffer}; +use arrow_buffer::{ArrowNativeType, MutableBuffer, NullBuffer}; use arrow_data::ArrayDataBuilder; use arrow_schema::{ArrowError, DataType}; @@ -168,6 +170,129 @@ pub fn concat_elements_utf8_many( Ok(unsafe { builder.build_unchecked() }.into()) } +/// Returns the elementwise concatenation of a [`FixedSizeBinaryArray`]. +/// +/// The result has `value_length = left.value_length() + right.value_length()`. +/// An index is null if either input is null at that position. +/// +/// An error will be returned if `left` and `right` have different lengths. +pub fn concat_elements_fixed_size_binary( + left: &FixedSizeBinaryArray, + right: &FixedSizeBinaryArray, +) -> Result { + if left.len() != right.len() { + return Err(ArrowError::ComputeError(format!( + "Arrays must have the same length: {} != {}", + left.len(), + right.len() + ))); + } + + let left_size = left.value_length() as usize; + let right_size = right.value_length() as usize; + let output_size = left_size + right_size; + + // Pre-compute combined null bitmap so the per-row NULL check is efficient + let nulls = NullBuffer::union(left.nulls(), right.nulls()); + + let mut result = FixedSizeBinaryBuilder::with_capacity(left.len(), output_size as i32); + let mut buffer = MutableBuffer::with_capacity(output_size); + for i in 0..left.len() { + if nulls.as_ref().is_some_and(|n| n.is_null(i)) { + result.append_null(); + } else { + buffer.clear(); + buffer.extend_from_slice(left.value(i)); + buffer.extend_from_slice(right.value(i)); + result.append_value(&buffer)?; + } + } + + Ok(result.finish()) +} + +/// Concatenates two `BinaryViewArray`s element-wise. +/// If either element is `Null`, the result element is also `Null`. +/// +/// # Errors +/// - Returns an error if the input arrays have different lengths. +/// - Returns an error if any concatenated value exceeds `u32::MAX` in length. +pub fn concat_elements_binary_view_array( + left: &BinaryViewArray, + right: &BinaryViewArray, +) -> Result { + if left.len() != right.len() { + return Err(ArrowError::ComputeError(format!( + "Arrays must have the same length: {} != {}", + left.len(), + right.len() + ))); + } + let mut result = BinaryViewBuilder::with_capacity(left.len()); + + // Avoid reallocations by writing to a reused buffer + let mut buffer = MutableBuffer::new(0); + + // Pre-compute combined null bitmap, so the per-row NULL check is efficient + let nulls = NullBuffer::union(left.nulls(), right.nulls()); + + for i in 0..left.len() { + if nulls.as_ref().is_some_and(|n| n.is_null(i)) { + result.append_null(); + } else { + buffer.clear(); + buffer.extend_from_slice(left.value(i)); + buffer.extend_from_slice(right.value(i)); + result.try_append_value(&buffer)?; + } + } + Ok(result.finish()) +} + +/// Concatenates two `StringViewArray`s element-wise. +/// If either element is `Null`, the result element is also `Null`. +/// +/// # Errors +/// - Returns an error if the input arrays have different lengths. +/// - Returns an error if any concatenated value exceeds `u32::MAX` in length. +/// - Returns an error if concatenated strings do not result in a proper UTF-8 string +// Cannot reuse code with `GenericByteViewBuilder` since `try_append_value` works with +// `AsRef`, and there is no conversion from `ByteViewType` to this or [u8] +pub fn concat_elements_string_view_array( + left: &StringViewArray, + right: &StringViewArray, +) -> Result { + if left.len() != right.len() { + return Err(ArrowError::ComputeError(format!( + "Arrays must have the same length: {} != {}", + left.len(), + right.len() + ))); + } + + let mut result = StringViewBuilder::with_capacity(left.len()); + + // Avoid reallocations by writing to a reused buffer + let mut buffer: Vec = Vec::new(); + + let nulls = NullBuffer::union(left.nulls(), right.nulls()); + + for i in 0..left.len() { + if nulls.as_ref().is_some_and(|n| n.is_null(i)) { + result.append_null(); + } else { + buffer.clear(); + buffer.extend_from_slice(left.value(i).as_bytes()); + buffer.extend_from_slice(right.value(i).as_bytes()); + let s = std::str::from_utf8(&buffer).map_err(|_| { + ArrowError::ComputeError("Concatenated values are not valid UTF-8".into()) + })?; + result.try_append_value(s)?; + } + } + Ok(result.finish()) +} + /// Returns the elementwise concatenation of [`Array`]s. /// /// # Errors @@ -185,22 +310,43 @@ pub fn concat_elements_dyn(left: &dyn Array, right: &dyn Array) -> Result { let left = left.as_any().downcast_ref::().unwrap(); let right = right.as_any().downcast_ref::().unwrap(); - Ok(Arc::new(concat_elements_utf8(left, right).unwrap())) + Ok(Arc::new(concat_elements_utf8(left, right)?)) } (DataType::LargeUtf8, DataType::LargeUtf8) => { let left = left.as_any().downcast_ref::().unwrap(); let right = right.as_any().downcast_ref::().unwrap(); - Ok(Arc::new(concat_elements_utf8(left, right).unwrap())) + Ok(Arc::new(concat_elements_utf8(left, right)?)) } (DataType::Binary, DataType::Binary) => { let left = left.as_any().downcast_ref::().unwrap(); let right = right.as_any().downcast_ref::().unwrap(); - Ok(Arc::new(concat_element_binary(left, right).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).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::FixedSizeBinary(_), DataType::FixedSizeBinary(_)) => { + let left = left + .as_any() + .downcast_ref::() + .unwrap(); + let right = right + .as_any() + .downcast_ref::() + .unwrap(); + Ok(Arc::new(concat_elements_fixed_size_binary(left, right)?)) } // unimplemented _ => Err(ArrowError::NotYetImplemented(format!( @@ -213,6 +359,8 @@ pub fn concat_elements_dyn(left: &dyn Array, right: &dyn Array) -> Result>); + let right = BinaryViewArray::from_iter(vec![] as Vec>); + + let output = concat_elements_binary_view_array(&left, &right).unwrap(); + let expected = BinaryViewArray::from_iter(vec![] as Vec>); + assert_eq!(output, expected); + } + #[test] fn test_concat_dyn_same_type() { // test for StringArray @@ -398,6 +661,26 @@ mod tests { .into(); let expected = LargeBinaryArray::from_opt_vec(vec![None, Some(b"baryyy"), None]); assert_eq!(output, expected); + + // test for BinaryViewArray + let left = BinaryViewArray::from_iter(vec![Some(b"foo" as &[u8]), Some(b"bar"), None]); + let right = BinaryViewArray::from_iter(vec![None, Some(b"yyy" as &[u8]), Some(b"zzz")]); + let output: BinaryViewArray = concat_elements_dyn(&left, &right) + .unwrap() + .into_data() + .into(); + let expected = BinaryViewArray::from_iter(vec![None, Some(b"baryyy" as &[u8]), None]); + assert_eq!(output, expected); + + // test for FixedSizeBinaryArray + let left = FixedSizeBinaryArray::from(vec![Some(b"foo" as &[u8]), Some(b"bar"), None]); + let right = FixedSizeBinaryArray::from(vec![None, Some(b"yyy" as &[u8]), Some(b"zzz")]); + let output: FixedSizeBinaryArray = concat_elements_dyn(&left, &right) + .unwrap() + .into_data() + .into(); + let expected = FixedSizeBinaryArray::from(vec![None, Some(b"baryyy" as &[u8]), None]); + assert_eq!(output, expected); } #[test]