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

arrow_string/
concat_elements.rs
