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

arrow_string/
concat_elements.rs
