arrow_string/
concat_elements.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Provides utility functions for concatenation of elements in arrays.
19
20use std::marker::PhantomData;
21use std::sync::Arc;
22
23use arrow_array::builder::{FixedSizeBinaryBuilder, make_view};
24use arrow_array::types::{ByteArrayType, ByteViewType};
25use arrow_array::*;
26use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer, NullBuffer, ScalarBuffer};
27use arrow_data::{ArrayDataBuilder, MAX_INLINE_VIEW_LEN};
28use arrow_schema::{ArrowError, DataType};
29
30/// Returns the elementwise concatenation of a [`GenericByteArray`].
31///
32/// # Errors
33///
34/// Returns an error if the arrays have different lengths, or if the concatenated
35/// data is too long for the offset type.
36pub fn concat_elements_bytes<T: ByteArrayType>(
37    left: &GenericByteArray<T>,
38    right: &GenericByteArray<T>,
39) -> Result<GenericByteArray<T>, ArrowError> {
40    if left.len() != right.len() {
41        return Err(ArrowError::ComputeError(format!(
42            "Arrays must have the same length: {} != {}",
43            left.len(),
44            right.len()
45        )));
46    }
47
48    let nulls = NullBuffer::union(left.nulls(), right.nulls());
49
50    let left_offsets = left.offsets();
51    let right_offsets = right.offsets();
52
53    let left_values = left.value_data();
54    let right_values = right.value_data();
55
56    let mut output_values = Vec::with_capacity(
57        (left_offsets.last() - left_offsets.first()).as_usize()
58            + (right_offsets.last() - right_offsets.first()).as_usize(),
59    );
60
61    let mut output_offsets = Vec::with_capacity(left_offsets.len());
62    output_offsets.push(T::Offset::usize_as(0));
63    for (left_idx, right_idx) in left_offsets.windows(2).zip(right_offsets.windows(2)) {
64        output_values
65            .extend_from_slice(&left_values[left_idx[0].as_usize()..left_idx[1].as_usize()]);
66        output_values
67            .extend_from_slice(&right_values[right_idx[0].as_usize()..right_idx[1].as_usize()]);
68        let output_len = output_values.len();
69        let offset =
70            T::Offset::from_usize(output_len).ok_or(ArrowError::OffsetOverflowError(output_len))?;
71        output_offsets.push(offset);
72    }
73
74    let builder = ArrayDataBuilder::new(T::DATA_TYPE)
75        .len(left.len())
76        .add_buffer(output_offsets.into())
77        .add_buffer(output_values.into())
78        .nulls(nulls);
79
80    // SAFETY - offsets valid by construction
81    Ok(unsafe { builder.build_unchecked() }.into())
82}
83
84/// Returns the elementwise concatenation of a [`GenericStringArray`].
85///
86/// An index of the resulting [`GenericStringArray`] is null if any of
87/// `StringArray` are null at that location.
88///
89/// ```text
90/// e.g:
91///
92///   ["Hello"] + ["World"] = ["HelloWorld"]
93///
94///   ["a", "b"] + [None, "c"] = [None, "bc"]
95/// ```
96///
97/// An error will be returned if `left` and `right` have different lengths
98pub fn concat_elements_utf8<Offset: OffsetSizeTrait>(
99    left: &GenericStringArray<Offset>,
100    right: &GenericStringArray<Offset>,
101) -> Result<GenericStringArray<Offset>, ArrowError> {
102    concat_elements_bytes(left, right)
103}
104
105/// Returns the elementwise concatenation of a [`GenericBinaryArray`].
106pub fn concat_element_binary<Offset: OffsetSizeTrait>(
107    left: &GenericBinaryArray<Offset>,
108    right: &GenericBinaryArray<Offset>,
109) -> Result<GenericBinaryArray<Offset>, ArrowError> {
110    concat_elements_bytes(left, right)
111}
112
113/// Returns the elementwise concatenation of [`StringArray`].
114/// ```text
115/// e.g:
116///   ["a", "b"] + [None, "c"] + [None, "d"] = [None, "bcd"]
117/// ```
118///
119/// An error will be returned if the [`StringArray`] are of different lengths
120///
121/// # Errors
122///
123/// Returns an error if the arrays have different lengths, or if the concatenated
124/// data is too long for the offset type.
125pub fn concat_elements_utf8_many<Offset: OffsetSizeTrait>(
126    arrays: &[&GenericStringArray<Offset>],
127) -> Result<GenericStringArray<Offset>, ArrowError> {
128    if arrays.is_empty() {
129        return Err(ArrowError::ComputeError(
130            "concat requires input of at least one array".to_string(),
131        ));
132    }
133
134    let size = arrays[0].len();
135    if !arrays.iter().all(|array| array.len() == size) {
136        return Err(ArrowError::ComputeError(format!(
137            "Arrays must have the same length of {size}",
138        )));
139    }
140
141    let nulls = arrays
142        .iter()
143        .fold(None, |acc, a| NullBuffer::union(acc.as_ref(), a.nulls()));
144
145    let data_values = arrays
146        .iter()
147        .map(|array| array.value_data())
148        .collect::<Vec<_>>();
149
150    let mut offsets = arrays
151        .iter()
152        .map(|a| a.value_offsets().iter().peekable())
153        .collect::<Vec<_>>();
154
155    let mut output_values = Vec::with_capacity(
156        arrays
157            .iter()
158            .map(|array| (array.offsets().last() - array.offsets().first()).as_usize())
159            .sum(),
160    );
161
162    let mut output_offsets = Vec::with_capacity(size + 1);
163    output_offsets.push(Offset::zero());
164    for _ in 0..size {
165        data_values
166            .iter()
167            .zip(offsets.iter_mut())
168            .for_each(|(values, offset)| {
169                let index_start = offset.next().unwrap().as_usize();
170                let index_end = offset.peek().unwrap().as_usize();
171                output_values.extend_from_slice(&values[index_start..index_end]);
172            });
173        let output_len = output_values.len();
174        let offset =
175            Offset::from_usize(output_len).ok_or(ArrowError::OffsetOverflowError(output_len))?;
176        output_offsets.push(offset);
177    }
178
179    let builder = ArrayDataBuilder::new(GenericStringArray::<Offset>::DATA_TYPE)
180        .len(size)
181        .add_buffer(output_offsets.into())
182        .add_buffer(output_values.into())
183        .nulls(nulls);
184
185    // SAFETY - offsets valid by construction
186    Ok(unsafe { builder.build_unchecked() }.into())
187}
188
189/// Returns the elementwise concatenation of a [`FixedSizeBinaryArray`].
190///
191/// The result has `value_length = left.value_length() + right.value_length()`.
192/// An index is null if either input is null at that position.
193///
194/// An error will be returned if `left` and `right` have different lengths.
195pub fn concat_elements_fixed_size_binary(
196    left: &FixedSizeBinaryArray,
197    right: &FixedSizeBinaryArray,
198) -> Result<FixedSizeBinaryArray, ArrowError> {
199    if left.len() != right.len() {
200        return Err(ArrowError::ComputeError(format!(
201            "Arrays must have the same length: {} != {}",
202            left.len(),
203            right.len()
204        )));
205    }
206
207    let left_size: usize = left.value_length().try_into().map_err(|_| {
208        ArrowError::InvalidArgumentError(format!(
209            "Invalid size of FixedSizeBinaryArray({})",
210            left.value_length()
211        ))
212    })?;
213    let right_size: usize = right.value_length().try_into().map_err(|_| {
214        ArrowError::InvalidArgumentError(format!(
215            "Invalid size of FixedSizeBinaryArray({})",
216            right.value_length()
217        ))
218    })?;
219    let output_size = left_size + right_size;
220    let output_value_length = i32::try_from(output_size).map_err(|_| {
221        ArrowError::InvalidArgumentError(format!(
222            "Concatenated FixedSizeBinary value length {output_size} exceeds i32"
223        ))
224    })?;
225
226    // Pre-compute combined null bitmap so the per-row NULL check is efficient
227    let nulls = NullBuffer::union(left.nulls(), right.nulls());
228
229    let mut result = FixedSizeBinaryBuilder::with_capacity(left.len(), output_value_length);
230    let mut buffer = MutableBuffer::with_capacity(output_size);
231    for i in 0..left.len() {
232        if nulls.as_ref().is_some_and(|n| n.is_null(i)) {
233            result.append_null();
234        } else {
235            buffer.clear();
236            buffer
237                .try_extend_from_slice(left.value(i))
238                .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
239            buffer
240                .try_extend_from_slice(right.value(i))
241                .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
242            result.append_value(&buffer)?;
243        }
244    }
245
246    Ok(result.finish())
247}
248
249struct ConcatByteViewBuilder<T>
250where
251    T: ByteViewType,
252{
253    views: Vec<u128>,
254    data: Vec<u8>,
255    inline: Vec<u8>,
256    phantom: PhantomData<T>,
257}
258
259impl<T> ConcatByteViewBuilder<T>
260where
261    T: ByteViewType,
262{
263    /// Returns the elementwise concatenation of two [`GenericByteViewArray`]s.
264    fn concat_elements_view_array(
265        left: &GenericByteViewArray<T>,
266        right: &GenericByteViewArray<T>,
267    ) -> Result<GenericByteViewArray<T>, ArrowError> {
268        let len = left.len();
269        if len != right.len() {
270            return Err(ArrowError::ComputeError(format!(
271                "Arrays must have the same length: {} != {}",
272                len,
273                right.len()
274            )));
275        }
276
277        let null_buffer = NullBuffer::union(left.nulls(), right.nulls());
278
279        // Compute the required data buffer size, excluding any elements that are null
280        // or are small enough to be stored inline. Both lengths are `u32`, and a
281        // pair of them can sum past that, so the addition is done in `usize`.
282        let inline_max = MAX_INLINE_VIEW_LEN as usize;
283        let data_size: usize = match &null_buffer {
284            None => left
285                .lengths()
286                .zip(right.lengths())
287                .map(|(l, r)| l as usize + r as usize)
288                .filter(|len| *len > inline_max)
289                .sum(),
290            Some(nb) => left
291                .lengths()
292                .zip(right.lengths())
293                .zip(nb.iter())
294                .filter(|((_, _), not_null)| *not_null)
295                .map(|((l, r), _)| l as usize + r as usize)
296                .filter(|len| *len > inline_max)
297                .sum(),
298        };
299
300        if data_size > i32::MAX as usize {
301            return Err(ArrowError::ArithmeticOverflow(
302                "byte array offset overflow".to_string(),
303            ));
304        }
305        let mut builder = Self::with_capacity(len, data_size);
306
307        match &null_buffer {
308            None => {
309                for (l, r) in left.bytes_iter().zip(right.bytes_iter()) {
310                    builder.append_concat_view(l, r);
311                }
312            }
313            Some(nb) => {
314                for ((l, r), not_null) in left.bytes_iter().zip(right.bytes_iter()).zip(nb.iter()) {
315                    if not_null {
316                        builder.append_concat_view(l, r);
317                    } else {
318                        builder.append_empty_view();
319                    }
320                }
321            }
322        }
323
324        builder.finish(null_buffer)
325    }
326
327    fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
328        Self {
329            views: Vec::with_capacity(item_capacity),
330            data: Vec::with_capacity(data_capacity),
331            inline: Vec::with_capacity(MAX_INLINE_VIEW_LEN as usize),
332            phantom: PhantomData,
333        }
334    }
335
336    /// Append a view containing the concatenation of `left` and `right`.
337    fn append_concat_view(&mut self, left: &[u8], right: &[u8]) {
338        let total_len = left.len() + right.len();
339        if total_len > MAX_INLINE_VIEW_LEN as usize {
340            let offset = self.data.len();
341
342            // SAFETY: we've checked that the total data size is within i32::MAX
343            // in `concat_elements_view_array`, so offset cannot exceed it.
344            // Not using `u32::try_from` on each insertion makes a ~5% difference
345            // in benchmarking
346            debug_assert!(i32::try_from(offset).is_ok());
347            let view_offset: u32 = offset as u32;
348
349            self.data.extend_from_slice(left);
350            self.data.extend_from_slice(right);
351            self.views
352                .push(make_view(&self.data[offset..], 0, view_offset));
353        } else {
354            self.inline.extend_from_slice(left);
355            self.inline.extend_from_slice(right);
356            self.views.push(make_view(&self.inline, 0, 0));
357            self.inline.clear();
358        }
359    }
360
361    /// Append an empty view.
362    #[inline]
363    fn append_empty_view(&mut self) {
364        self.views.push(0);
365    }
366
367    fn finish(
368        self,
369        null_buffer: Option<NullBuffer>,
370    ) -> Result<GenericByteViewArray<T>, ArrowError> {
371        if let Some(ref nulls) = null_buffer
372            && nulls.len() != self.views.len()
373        {
374            return Err(ArrowError::ComputeError(format!(
375                "Null buffer length ({}) must match row count ({})",
376                nulls.len(),
377                self.views.len()
378            )));
379        }
380
381        let buffers = if self.data.is_empty() {
382            Arc::from([])
383        } else {
384            Arc::from([Buffer::from(self.data)])
385        };
386
387        // SAFETY: views were constructed with correct lengths, offsets, and
388        // prefixes. UTF-8 validity is implicitly guaranteed by never concatenating
389        // arrays with mixed ByteViewTypes.
390        let array = unsafe {
391            GenericByteViewArray::<T>::new_unchecked(
392                ScalarBuffer::from(self.views),
393                buffers,
394                null_buffer,
395            )
396        };
397        Ok(array)
398    }
399}
400
401/// Concatenates two `BinaryViewArray`s element-wise.
402/// If either element is `Null`, the result element is also `Null`.
403///
404/// # Errors
405/// - Returns an error if the input arrays have different lengths.
406/// - Returns an error if any concatenated value exceeds `u32::MAX` in length.
407pub fn concat_elements_binary_view_array(
408    left: &BinaryViewArray,
409    right: &BinaryViewArray,
410) -> Result<BinaryViewArray, ArrowError> {
411    ConcatByteViewBuilder::concat_elements_view_array(left, right)
412}
413
414/// Concatenates two `StringViewArray`s element-wise.
415/// If either element is `Null`, the result element is also `Null`.
416///
417/// # Errors
418/// - Returns an error if the input arrays have different lengths.
419/// - Returns an error if any concatenated value exceeds `u32::MAX` in length.
420/// - Returns an error if concatenated strings do not result in a proper UTF-8 string
421pub fn concat_elements_string_view_array(
422    left: &StringViewArray,
423    right: &StringViewArray,
424) -> Result<StringViewArray, ArrowError> {
425    ConcatByteViewBuilder::concat_elements_view_array(left, right)
426}
427
428/// Returns the elementwise concatenation of [`Array`]s.
429///
430/// The output array will have the same type as the input arrays (which must have the same type).
431///
432/// Concatenation of `FixedSizeBinaryArray` instances with different sizes is supported. The output
433/// type is `FixedSizeBinaryArray` with the sum of the sizes of the two input arrays as size.
434///
435/// # Errors
436///
437/// This function errors if the arrays are of different types.
438pub fn concat_elements_dyn(left: &dyn Array, right: &dyn Array) -> Result<ArrayRef, ArrowError> {
439    match (left.data_type(), right.data_type()) {
440        (DataType::Utf8, DataType::Utf8) => {
441            let left = left.as_any().downcast_ref::<StringArray>().unwrap();
442            let right = right.as_any().downcast_ref::<StringArray>().unwrap();
443            Ok(Arc::new(concat_elements_utf8(left, right)?))
444        }
445        (DataType::Utf8View, DataType::Utf8View) => {
446            let left = left.as_any().downcast_ref::<StringViewArray>().unwrap();
447            let right = right.as_any().downcast_ref::<StringViewArray>().unwrap();
448            Ok(Arc::new(concat_elements_string_view_array(left, right)?))
449        }
450        (DataType::LargeUtf8, DataType::LargeUtf8) => {
451            let left = left.as_any().downcast_ref::<LargeStringArray>().unwrap();
452            let right = right.as_any().downcast_ref::<LargeStringArray>().unwrap();
453            Ok(Arc::new(concat_elements_utf8(left, right)?))
454        }
455        (DataType::Binary, DataType::Binary) => {
456            let left = left.as_any().downcast_ref::<BinaryArray>().unwrap();
457            let right = right.as_any().downcast_ref::<BinaryArray>().unwrap();
458            Ok(Arc::new(concat_element_binary(left, right)?))
459        }
460        (DataType::BinaryView, DataType::BinaryView) => {
461            let left = left.as_any().downcast_ref::<BinaryViewArray>().unwrap();
462            let right = right.as_any().downcast_ref::<BinaryViewArray>().unwrap();
463            Ok(Arc::new(concat_elements_binary_view_array(left, right)?))
464        }
465        (DataType::LargeBinary, DataType::LargeBinary) => {
466            let left = left.as_any().downcast_ref::<LargeBinaryArray>().unwrap();
467            let right = right.as_any().downcast_ref::<LargeBinaryArray>().unwrap();
468            Ok(Arc::new(concat_element_binary(left, right)?))
469        }
470        (DataType::FixedSizeBinary(_), DataType::FixedSizeBinary(_)) => {
471            let left = left
472                .as_any()
473                .downcast_ref::<FixedSizeBinaryArray>()
474                .unwrap();
475            let right = right
476                .as_any()
477                .downcast_ref::<FixedSizeBinaryArray>()
478                .unwrap();
479            Ok(Arc::new(concat_elements_fixed_size_binary(left, right)?))
480        }
481        // unimplemented
482        (l, r) => {
483            if l != r {
484                Err(ArrowError::ComputeError(format!(
485                    "Cannot concat arrays of different types: {l} != {r}"
486                )))
487            } else {
488                Err(ArrowError::NotYetImplemented(format!(
489                    "concat not supported for {}",
490                    left.data_type()
491                )))
492            }
493        }
494    }
495}
496
497#[cfg(test)]
498mod tests {
499    use super::*;
500    use arrow_array::cast::AsArray;
501    use arrow_buffer::Buffer;
502
503    #[test]
504    fn test_string_concat() {
505        let left = [Some("foo"), Some("bar"), None]
506            .into_iter()
507            .collect::<StringArray>();
508        let right = [None, Some("yyy"), Some("zzz")]
509            .into_iter()
510            .collect::<StringArray>();
511
512        let output = concat_elements_dyn(&left, &right).unwrap();
513
514        let expected = [None, Some("baryyy"), None]
515            .into_iter()
516            .collect::<StringArray>();
517
518        assert_eq!(output.as_string(), &expected);
519    }
520
521    #[test]
522    fn test_string_concat_empty_string() {
523        let left = [Some("foo"), Some(""), Some("bar")]
524            .into_iter()
525            .collect::<StringArray>();
526        let right = [Some("baz"), Some(""), Some("")]
527            .into_iter()
528            .collect::<StringArray>();
529
530        let output = concat_elements_dyn(&left, &right).unwrap();
531
532        let expected = [Some("foobaz"), Some(""), Some("bar")]
533            .into_iter()
534            .collect::<StringArray>();
535
536        assert_eq!(output.as_string(), &expected);
537    }
538
539    #[test]
540    fn test_string_concat_no_null() {
541        let left = StringArray::from(vec!["foo", "bar"]);
542        let right = StringArray::from(vec!["bar", "baz"]);
543
544        let output = concat_elements_dyn(&left, &right).unwrap();
545
546        let expected = StringArray::from(vec!["foobar", "barbaz"]);
547
548        assert_eq!(output.as_string(), &expected);
549    }
550
551    #[test]
552    fn test_string_concat_error() {
553        let left = StringArray::from(vec!["foo", "bar"]);
554        let right = StringArray::from(vec!["baz"]);
555
556        let output = concat_elements_dyn(&left, &right);
557
558        assert_eq!(
559            output.unwrap_err().to_string(),
560            "Compute error: Arrays must have the same length: 2 != 1".to_string()
561        );
562    }
563
564    #[test]
565    fn test_string_concat_slice() {
566        let left = &StringArray::from(vec![None, Some("foo"), Some("bar"), Some("baz")]);
567        let right = &StringArray::from(vec![Some("boo"), None, Some("far"), Some("faz")]);
568
569        let left_slice = left.slice(0, 3);
570        let right_slice = right.slice(1, 3);
571        let output = concat_elements_dyn(
572            left_slice
573                .as_any()
574                .downcast_ref::<GenericStringArray<i32>>()
575                .unwrap(),
576            right_slice
577                .as_any()
578                .downcast_ref::<GenericStringArray<i32>>()
579                .unwrap(),
580        )
581        .unwrap();
582
583        let expected = [None, Some("foofar"), Some("barfaz")]
584            .into_iter()
585            .collect::<StringArray>();
586
587        assert_eq!(output.as_string(), &expected);
588
589        let left_slice = left.slice(2, 2);
590        let right_slice = right.slice(1, 2);
591
592        let output = concat_elements_dyn(
593            left_slice
594                .as_any()
595                .downcast_ref::<GenericStringArray<i32>>()
596                .unwrap(),
597            right_slice
598                .as_any()
599                .downcast_ref::<GenericStringArray<i32>>()
600                .unwrap(),
601        )
602        .unwrap();
603
604        let expected = [None, Some("bazfar")].into_iter().collect::<StringArray>();
605
606        assert_eq!(output.as_string(), &expected);
607    }
608
609    #[test]
610    fn test_concat_slice_capacity() {
611        let left = StringArray::from(vec!["hello"; 8]).slice(2, 3);
612        let right = StringArray::from(vec![" world"; 8]).slice(4, 3);
613
614        let pair = concat_elements_utf8(&left, &right).unwrap();
615        assert_eq!(pair, StringArray::from(vec!["hello world"; 3]));
616        assert_eq!(pair.values().capacity(), 33);
617
618        let many = concat_elements_utf8_many(&[&left, &right, &left]).unwrap();
619        assert_eq!(many, StringArray::from(vec!["hello worldhello"; 3]));
620        assert_eq!(many.values().capacity(), 48);
621    }
622
623    #[test]
624    fn test_string_concat_error_empty() {
625        assert_eq!(
626            concat_elements_utf8_many::<i32>(&[])
627                .unwrap_err()
628                .to_string(),
629            "Compute error: concat requires input of at least one array".to_string()
630        );
631    }
632
633    #[test]
634    fn test_string_concat_one() {
635        let expected = [None, Some("baryyy"), None]
636            .into_iter()
637            .collect::<StringArray>();
638
639        let output = concat_elements_utf8_many(&[&expected]).unwrap();
640
641        assert_eq!(output, expected);
642    }
643
644    #[test]
645    fn test_string_concat_many() {
646        let foo = StringArray::from(vec![Some("f"), Some("o"), Some("o"), None]);
647        let bar = StringArray::from(vec![None, Some("b"), Some("a"), Some("r")]);
648        let baz = StringArray::from(vec![Some("b"), None, Some("a"), Some("z")]);
649
650        let output = concat_elements_utf8_many(&[&foo, &bar, &baz]).unwrap();
651
652        let expected = [None, None, Some("oaa"), None]
653            .into_iter()
654            .collect::<StringArray>();
655
656        assert_eq!(output, expected);
657    }
658
659    #[test]
660    fn test_fixed_size_binary_concat() {
661        let left = FixedSizeBinaryArray::try_from(vec![Some(b"foo" as &[u8]), Some(b"bar"), None])
662            .unwrap();
663        let right = FixedSizeBinaryArray::try_from(vec![None, Some(b"yyy" as &[u8]), Some(b"zzz")])
664            .unwrap();
665
666        let output = concat_elements_dyn(&left, &right).unwrap();
667
668        let expected =
669            FixedSizeBinaryArray::try_from(vec![None, Some(b"baryyy" as &[u8]), None]).unwrap();
670        assert_eq!(output.as_fixed_size_binary(), &expected);
671    }
672
673    #[test]
674    fn test_mixed_fixed_size_binary_concat() {
675        let left =
676            FixedSizeBinaryArray::try_from(vec![Some(b"foobar" as &[u8]), Some(b"barbaz"), None])
677                .unwrap();
678        let right = FixedSizeBinaryArray::try_from(vec![None, Some(b"yyy" as &[u8]), Some(b"zzz")])
679            .unwrap();
680
681        let output = concat_elements_dyn(&left, &right).unwrap();
682
683        let expected =
684            FixedSizeBinaryArray::try_from(vec![None, Some(b"barbazyyy" as &[u8]), None]).unwrap();
685        assert_eq!(output.as_fixed_size_binary(), &expected);
686    }
687
688    #[test]
689    fn test_fixed_size_binary_concat_no_null() {
690        let left = FixedSizeBinaryArray::try_from(vec![b"ab" as &[u8], b"cd"]).unwrap();
691        let right = FixedSizeBinaryArray::try_from(vec![b"12" as &[u8], b"34"]).unwrap();
692
693        let output = concat_elements_dyn(&left, &right).unwrap();
694
695        let expected = FixedSizeBinaryArray::try_from(vec![b"ab12" as &[u8], b"cd34"]).unwrap();
696        assert_eq!(output.as_fixed_size_binary(), &expected);
697    }
698
699    #[test]
700    fn test_fixed_size_binary_concat_error() {
701        let left = FixedSizeBinaryArray::try_from(vec![b"ab" as &[u8], b"cd"]).unwrap();
702        let right = FixedSizeBinaryArray::try_from(vec![b"12" as &[u8]]).unwrap();
703
704        let output = concat_elements_dyn(&left, &right);
705        assert_eq!(
706            output.unwrap_err().to_string(),
707            "Compute error: Arrays must have the same length: 2 != 1".to_string()
708        );
709    }
710
711    #[test]
712    fn test_fixed_size_binary_concat_width_overflow() {
713        let width = 0x7000_0000_i32;
714        let left =
715            FixedSizeBinaryArray::try_new(width, Buffer::from(Vec::<u8>::new()), None).unwrap();
716        let right =
717            FixedSizeBinaryArray::try_new(width, Buffer::from(Vec::<u8>::new()), None).unwrap();
718
719        let output = concat_elements_fixed_size_binary(&left, &right);
720        assert_eq!(
721            output.unwrap_err().to_string(),
722            "Invalid argument error: Concatenated FixedSizeBinary value length 3758096384 exceeds i32".to_string()
723        );
724    }
725
726    #[test]
727    fn test_fixed_size_binary_concat_empty() {
728        let left = FixedSizeBinaryArray::new(0, Buffer::from(&[]), None);
729        let right = FixedSizeBinaryArray::new(0, Buffer::from(&[]), None);
730
731        let output = concat_elements_dyn(&left, &right).unwrap();
732
733        let expected = FixedSizeBinaryArray::new(0, Buffer::from(&[]), None);
734        assert_eq!(output.as_fixed_size_binary(), &expected);
735    }
736
737    #[test]
738    fn test_binary_view_concat() {
739        let long = b"ThisStringIsLongerThan12Bytes" as &[u8];
740        let left = BinaryViewArray::from_iter(vec![
741            Some(b"foo" as &[u8]),
742            Some(b"bar"),
743            None,
744            Some(b"foofoofoo"),
745            Some(b"foo"),
746            Some(long),
747            Some(long),
748        ]);
749        let right = BinaryViewArray::from_iter(vec![
750            None,
751            Some(b"yyy" as &[u8]),
752            Some(b"zzz"),
753            Some(b"barbarbar"),
754            Some(long),
755            Some(b"bar"),
756            Some(long),
757        ]);
758
759        let output = concat_elements_binary_view_array(&left, &right).unwrap();
760
761        let expected = BinaryViewArray::from_iter(vec![
762            None,
763            Some(b"baryyy" as &[u8]),
764            None,
765            Some(b"foofoofoobarbarbar"),
766            Some(b"fooThisStringIsLongerThan12Bytes"),
767            Some(b"ThisStringIsLongerThan12Bytesbar"),
768            Some(b"ThisStringIsLongerThan12BytesThisStringIsLongerThan12Bytes"),
769        ]);
770        assert_eq!(output, expected);
771    }
772
773    #[test]
774    fn test_string_view_concat() {
775        let long = "ThisStringIsLongerThan12Bytes";
776        let left = StringViewArray::from_iter(vec![
777            Some("foo"),
778            Some("bar"),
779            None,
780            Some("foofoofoo"),
781            Some("foo"),
782            Some(long),
783            Some(long),
784        ]);
785        let right = StringViewArray::from_iter(vec![
786            None,
787            Some("yyy"),
788            Some("zzz"),
789            Some("barbarbar"),
790            Some(long),
791            Some("bar"),
792            Some(long),
793        ]);
794
795        let output = concat_elements_dyn(&left, &right).unwrap();
796
797        let expected = StringViewArray::from_iter(vec![
798            None,
799            Some("baryyy"),
800            None,
801            Some("foofoofoobarbarbar"),
802            Some("fooThisStringIsLongerThan12Bytes"),
803            Some("ThisStringIsLongerThan12Bytesbar"),
804            Some("ThisStringIsLongerThan12BytesThisStringIsLongerThan12Bytes"),
805        ]);
806        assert_eq!(output.as_string_view(), &expected);
807
808        let left = StringViewArray::from_iter(vec![
809            Some("a"),
810            Some("b"),
811            Some("foofoofoo"),
812            Some("a"),
813            Some(long),
814            Some(long),
815        ]);
816        let right = StringViewArray::from_iter(vec![
817            Some("c"),
818            Some("d"),
819            Some("barbarbar"),
820            Some(long),
821            Some("d"),
822            Some(long),
823        ]);
824
825        let output = concat_elements_dyn(&left, &right).unwrap();
826
827        let expected = StringViewArray::from_iter(vec![
828            Some("ac"),
829            Some("bd"),
830            Some("foofoofoobarbarbar"),
831            Some("aThisStringIsLongerThan12Bytes"),
832            Some("ThisStringIsLongerThan12Bytesd"),
833            Some("ThisStringIsLongerThan12BytesThisStringIsLongerThan12Bytes"),
834        ]);
835        assert_eq!(output.as_string_view(), &expected);
836    }
837
838    #[test]
839    fn test_binary_view_concat_no_null() {
840        let left = BinaryViewArray::from_iter(vec![
841            Some(b"foo" as &[u8]),
842            Some(b"bar"),
843            Some(b""),
844            Some(b"baz"),
845        ]);
846        let right = BinaryViewArray::from_iter(vec![
847            Some(b"bar" as &[u8]),
848            Some(b"baz"),
849            Some(b""),
850            Some(b""),
851        ]);
852
853        let output = concat_elements_dyn(&left, &right).unwrap();
854
855        let expected = BinaryViewArray::from_iter(vec![
856            Some(b"foobar" as &[u8]),
857            Some(b"barbaz"),
858            Some(b""),
859            Some(b"baz"),
860        ]);
861        assert_eq!(output.as_binary_view(), &expected);
862    }
863
864    #[test]
865    fn test_binary_view_concat_error() {
866        let left = BinaryViewArray::from_iter(vec![Some(b"foo" as &[u8]), Some(b"bar")]);
867        let right = BinaryViewArray::from_iter(vec![Some(b"baz" as &[u8])]);
868
869        let output = concat_elements_dyn(&left, &right);
870        assert_eq!(
871            output.unwrap_err().to_string(),
872            "Compute error: Arrays must have the same length: 2 != 1".to_string()
873        );
874    }
875
876    #[test]
877    fn test_binary_view_concat_empty() {
878        let left = BinaryViewArray::from_iter(vec![] as Vec<Option<&[u8]>>);
879        let right = BinaryViewArray::from_iter(vec![] as Vec<Option<&[u8]>>);
880
881        let output = concat_elements_dyn(&left, &right).unwrap();
882        let expected = BinaryViewArray::from_iter(vec![] as Vec<Option<&[u8]>>);
883        assert_eq!(output.as_binary_view(), &expected);
884    }
885
886    #[test]
887    fn test_concat_dyn_same_type() {
888        // test for StringArray
889        let left = StringArray::from(vec![Some("foo"), Some("bar"), None]);
890        let right = StringArray::from(vec![None, Some("yyy"), Some("zzz")]);
891
892        let output = concat_elements_dyn(&left, &right).unwrap();
893        let expected = StringArray::from(vec![None, Some("baryyy"), None]);
894        assert_eq!(output.as_string(), &expected);
895
896        // test for LargeStringArray
897        let left = LargeStringArray::from(vec![Some("foo"), Some("bar"), None]);
898        let right = LargeStringArray::from(vec![None, Some("yyy"), Some("zzz")]);
899
900        let output = concat_elements_dyn(&left, &right).unwrap();
901        let expected = LargeStringArray::from(vec![None, Some("baryyy"), None]);
902        assert_eq!(output.as_string(), &expected);
903
904        // test for BinaryArray
905        let left = BinaryArray::from_opt_vec(vec![Some(b"foo"), Some(b"bar"), None]);
906        let right = BinaryArray::from_opt_vec(vec![None, Some(b"yyy"), Some(b"zzz")]);
907        let output = concat_elements_dyn(&left, &right).unwrap();
908        let expected = BinaryArray::from_opt_vec(vec![None, Some(b"baryyy"), None]);
909        assert_eq!(output.as_binary(), &expected);
910
911        // test for LargeBinaryArray
912        let left = LargeBinaryArray::from_opt_vec(vec![Some(b"foo"), Some(b"bar"), None]);
913        let right = LargeBinaryArray::from_opt_vec(vec![None, Some(b"yyy"), Some(b"zzz")]);
914        let output = concat_elements_dyn(&left, &right).unwrap();
915        let expected = LargeBinaryArray::from_opt_vec(vec![None, Some(b"baryyy"), None]);
916        assert_eq!(output.as_binary(), &expected);
917
918        // test for BinaryViewArray
919        let long = b"ThisStringIsLongerThan12Bytes" as &[u8];
920        let left = BinaryViewArray::from_iter(vec![
921            Some(b"foo" as &[u8]),
922            Some(b"bar"),
923            None,
924            Some(b"foofoofoo"),
925            Some(b"foo"),
926            Some(long),
927            Some(long),
928        ]);
929        let right = BinaryViewArray::from_iter(vec![
930            None,
931            Some(b"yyy" as &[u8]),
932            Some(b"zzz"),
933            Some(b"barbarbar"),
934            Some(long),
935            Some(b"bar"),
936            Some(long),
937        ]);
938        let output = concat_elements_dyn(&left, &right).unwrap();
939        let expected = BinaryViewArray::from_iter(vec![
940            None,
941            Some(b"baryyy" as &[u8]),
942            None,
943            Some(b"foofoofoobarbarbar"),
944            Some(b"fooThisStringIsLongerThan12Bytes"),
945            Some(b"ThisStringIsLongerThan12Bytesbar"),
946            Some(b"ThisStringIsLongerThan12BytesThisStringIsLongerThan12Bytes"),
947        ]);
948        assert_eq!(output.as_binary_view(), &expected);
949
950        // test for StringViewArray
951        let long = "ThisStringIsLongerThan12Bytes";
952        let left = StringViewArray::from_iter(vec![
953            Some("foo"),
954            Some("bar"),
955            None,
956            Some("foofoofoo"),
957            Some("foo"),
958            Some(long),
959            Some(long),
960        ]);
961        let right = StringViewArray::from_iter(vec![
962            None,
963            Some("yyy"),
964            Some("zzz"),
965            Some("barbarbar"),
966            Some(long),
967            Some("bar"),
968            Some(long),
969        ]);
970        let output = concat_elements_dyn(&left, &right).unwrap();
971        let expected = StringViewArray::from_iter(vec![
972            None,
973            Some("baryyy"),
974            None,
975            Some("foofoofoobarbarbar"),
976            Some("fooThisStringIsLongerThan12Bytes"),
977            Some("ThisStringIsLongerThan12Bytesbar"),
978            Some("ThisStringIsLongerThan12BytesThisStringIsLongerThan12Bytes"),
979        ]);
980        assert_eq!(output.as_string_view(), &expected);
981
982        // test for FixedSizeBinaryArray
983        let left = FixedSizeBinaryArray::try_from(vec![Some(b"foo" as &[u8]), Some(b"bar"), None])
984            .unwrap();
985        let right = FixedSizeBinaryArray::try_from(vec![None, Some(b"yyy" as &[u8]), Some(b"zzz")])
986            .unwrap();
987        let output = concat_elements_dyn(&left, &right).unwrap();
988        let expected =
989            FixedSizeBinaryArray::try_from(vec![None, Some(b"baryyy" as &[u8]), None]).unwrap();
990        assert_eq!(output.as_fixed_size_binary(), &expected);
991    }
992
993    #[test]
994    fn test_concat_dyn_different_type() {
995        let left = StringArray::from(vec![Some("foo"), Some("bar"), None]);
996        let right = LargeStringArray::from(vec![None, Some("1"), Some("2")]);
997
998        let output = concat_elements_dyn(&left, &right);
999        assert_eq!(
1000            output.unwrap_err().to_string(),
1001            "Compute error: Cannot concat arrays of different types: Utf8 != LargeUtf8".to_string()
1002        );
1003    }
1004}
arrow_string/concat_elements.rs

arrow_string/
concat_elements.rs
