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

arrow_string/
concat_elements.rs
