arrow_select/
concat.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//! Defines concat kernel for `ArrayRef`
19//!
20//! Example:
21//!
22//! ```
23//! use arrow_array::{ArrayRef, StringArray};
24//! use arrow_select::concat::concat;
25//!
26//! let arr = concat(&[
27//!     &StringArray::from(vec!["hello", "world"]),
28//!     &StringArray::from(vec!["!"]),
29//! ]).unwrap();
30//! assert_eq!(arr.len(), 3);
31//! ```
32
33use crate::dictionary::{merge_dictionary_values, should_merge_dictionary_values};
34use arrow_array::builder::{
35    BooleanBuilder, GenericByteBuilder, GenericByteViewBuilder, PrimitiveBuilder,
36};
37use arrow_array::cast::AsArray;
38use arrow_array::types::*;
39use arrow_array::*;
40use arrow_buffer::{ArrowNativeType, BooleanBufferBuilder, NullBuffer, OffsetBuffer};
41use arrow_data::transform::{Capacities, MutableArrayData};
42use arrow_data::ArrayDataBuilder;
43use arrow_schema::{ArrowError, DataType, FieldRef, Fields, SchemaRef};
44use std::{collections::HashSet, ops::Add, sync::Arc};
45
46fn binary_capacity<T: ByteArrayType>(arrays: &[&dyn Array]) -> Capacities {
47    let mut item_capacity = 0;
48    let mut bytes_capacity = 0;
49    for array in arrays {
50        let a = array.as_bytes::<T>();
51
52        // Guaranteed to always have at least one element
53        let offsets = a.value_offsets();
54        bytes_capacity += offsets[offsets.len() - 1].as_usize() - offsets[0].as_usize();
55        item_capacity += a.len()
56    }
57
58    Capacities::Binary(item_capacity, Some(bytes_capacity))
59}
60
61fn fixed_size_list_capacity(arrays: &[&dyn Array], data_type: &DataType) -> Capacities {
62    if let DataType::FixedSizeList(f, _) = data_type {
63        let item_capacity = arrays.iter().map(|a| a.len()).sum();
64        let child_data_type = f.data_type();
65        match child_data_type {
66            // These types should match the types that `get_capacity`
67            // has special handling for.
68            DataType::Utf8
69            | DataType::LargeUtf8
70            | DataType::Binary
71            | DataType::LargeBinary
72            | DataType::FixedSizeList(_, _) => {
73                let values: Vec<&dyn arrow_array::Array> = arrays
74                    .iter()
75                    .map(|a| a.as_fixed_size_list().values().as_ref())
76                    .collect();
77                Capacities::List(
78                    item_capacity,
79                    Some(Box::new(get_capacity(&values, child_data_type))),
80                )
81            }
82            _ => Capacities::Array(item_capacity),
83        }
84    } else {
85        unreachable!("illegal data type for fixed size list")
86    }
87}
88
89fn concat_byte_view<B: ByteViewType>(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError> {
90    let mut builder =
91        GenericByteViewBuilder::<B>::with_capacity(arrays.iter().map(|a| a.len()).sum());
92    for &array in arrays.iter() {
93        builder.append_array(array.as_byte_view());
94    }
95    Ok(Arc::new(builder.finish()))
96}
97
98fn concat_dictionaries<K: ArrowDictionaryKeyType>(
99    arrays: &[&dyn Array],
100) -> Result<ArrayRef, ArrowError> {
101    let mut output_len = 0;
102    let dictionaries: Vec<_> = arrays
103        .iter()
104        .map(|x| x.as_dictionary::<K>())
105        .inspect(|d| output_len += d.len())
106        .collect();
107
108    if !should_merge_dictionary_values::<K>(&dictionaries, output_len) {
109        return concat_fallback(arrays, Capacities::Array(output_len));
110    }
111
112    let merged = merge_dictionary_values(&dictionaries, None)?;
113
114    // Recompute keys
115    let mut key_values = Vec::with_capacity(output_len);
116
117    let mut has_nulls = false;
118    for (d, mapping) in dictionaries.iter().zip(merged.key_mappings) {
119        has_nulls |= d.null_count() != 0;
120        for key in d.keys().values() {
121            // Use get to safely handle nulls
122            key_values.push(mapping.get(key.as_usize()).copied().unwrap_or_default())
123        }
124    }
125
126    let nulls = has_nulls.then(|| {
127        let mut nulls = BooleanBufferBuilder::new(output_len);
128        for d in &dictionaries {
129            match d.nulls() {
130                Some(n) => nulls.append_buffer(n.inner()),
131                None => nulls.append_n(d.len(), true),
132            }
133        }
134        NullBuffer::new(nulls.finish())
135    });
136
137    let keys = PrimitiveArray::<K>::new(key_values.into(), nulls);
138    // Sanity check
139    assert_eq!(keys.len(), output_len);
140
141    let array = unsafe { DictionaryArray::new_unchecked(keys, merged.values) };
142    Ok(Arc::new(array))
143}
144
145fn concat_lists<OffsetSize: OffsetSizeTrait>(
146    arrays: &[&dyn Array],
147    field: &FieldRef,
148) -> Result<ArrayRef, ArrowError> {
149    let mut output_len = 0;
150    let mut list_has_nulls = false;
151    let mut list_has_slices = false;
152
153    let lists = arrays
154        .iter()
155        .map(|x| x.as_list::<OffsetSize>())
156        .inspect(|l| {
157            output_len += l.len();
158            list_has_nulls |= l.null_count() != 0;
159            list_has_slices |= l.offsets()[0] > OffsetSize::zero()
160                || l.offsets().last().unwrap().as_usize() < l.values().len();
161        })
162        .collect::<Vec<_>>();
163
164    let lists_nulls = list_has_nulls.then(|| {
165        let mut nulls = BooleanBufferBuilder::new(output_len);
166        for l in &lists {
167            match l.nulls() {
168                Some(n) => nulls.append_buffer(n.inner()),
169                None => nulls.append_n(l.len(), true),
170            }
171        }
172        NullBuffer::new(nulls.finish())
173    });
174
175    // If any of the lists have slices, we need to slice the values
176    // to ensure that the offsets are correct
177    let mut sliced_values;
178    let values: Vec<&dyn Array> = if list_has_slices {
179        sliced_values = Vec::with_capacity(lists.len());
180        for l in &lists {
181            // if the first offset is non-zero, we need to slice the values so when
182            // we concatenate them below only the relevant values are included
183            let offsets = l.offsets();
184            let start_offset = offsets[0].as_usize();
185            let end_offset = offsets.last().unwrap().as_usize();
186            sliced_values.push(l.values().slice(start_offset, end_offset - start_offset));
187        }
188        sliced_values.iter().map(|a| a.as_ref()).collect()
189    } else {
190        lists.iter().map(|x| x.values().as_ref()).collect()
191    };
192
193    let concatenated_values = concat(values.as_slice())?;
194
195    // Merge value offsets from the lists
196    let value_offset_buffer =
197        OffsetBuffer::<OffsetSize>::from_lengths(lists.iter().flat_map(|x| x.offsets().lengths()));
198
199    let array = GenericListArray::<OffsetSize>::try_new(
200        Arc::clone(field),
201        value_offset_buffer,
202        concatenated_values,
203        lists_nulls,
204    )?;
205
206    Ok(Arc::new(array))
207}
208
209fn concat_primitives<T: ArrowPrimitiveType>(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError> {
210    let mut builder = PrimitiveBuilder::<T>::with_capacity(arrays.iter().map(|a| a.len()).sum())
211        .with_data_type(arrays[0].data_type().clone());
212
213    for array in arrays {
214        builder.append_array(array.as_primitive());
215    }
216
217    Ok(Arc::new(builder.finish()))
218}
219
220fn concat_boolean(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError> {
221    let mut builder = BooleanBuilder::with_capacity(arrays.iter().map(|a| a.len()).sum());
222
223    for array in arrays {
224        builder.append_array(array.as_boolean());
225    }
226
227    Ok(Arc::new(builder.finish()))
228}
229
230fn concat_bytes<T: ByteArrayType>(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError> {
231    let (item_capacity, bytes_capacity) = match binary_capacity::<T>(arrays) {
232        Capacities::Binary(item_capacity, Some(bytes_capacity)) => (item_capacity, bytes_capacity),
233        _ => unreachable!(),
234    };
235
236    let mut builder = GenericByteBuilder::<T>::with_capacity(item_capacity, bytes_capacity);
237
238    for array in arrays {
239        builder.append_array(array.as_bytes::<T>());
240    }
241
242    Ok(Arc::new(builder.finish()))
243}
244
245fn concat_structs(arrays: &[&dyn Array], fields: &Fields) -> Result<ArrayRef, ArrowError> {
246    let mut len = 0;
247    let mut has_nulls = false;
248    let structs = arrays
249        .iter()
250        .map(|a| {
251            len += a.len();
252            has_nulls |= a.null_count() > 0;
253            a.as_struct()
254        })
255        .collect::<Vec<_>>();
256
257    let nulls = has_nulls.then(|| {
258        let mut b = BooleanBufferBuilder::new(len);
259        for s in &structs {
260            match s.nulls() {
261                Some(n) => b.append_buffer(n.inner()),
262                None => b.append_n(s.len(), true),
263            }
264        }
265        NullBuffer::new(b.finish())
266    });
267
268    let column_concat_result = (0..fields.len())
269        .map(|i| {
270            let extracted_cols = structs
271                .iter()
272                .map(|s| s.column(i).as_ref())
273                .collect::<Vec<_>>();
274            concat(&extracted_cols)
275        })
276        .collect::<Result<Vec<_>, ArrowError>>()?;
277
278    Ok(Arc::new(StructArray::try_new(
279        fields.clone(),
280        column_concat_result,
281        nulls,
282    )?))
283}
284
285/// Concatenate multiple RunArray instances into a single RunArray.
286///
287/// This function handles the special case of concatenating RunArrays by:
288/// 1. Collecting all run ends and values from input arrays
289/// 2. Adjusting run ends to account for the length of previous arrays
290/// 3. Creating a new RunArray with the combined data
291fn concat_run_arrays<R: RunEndIndexType>(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError>
292where
293    R::Native: Add<Output = R::Native>,
294{
295    let run_arrays: Vec<_> = arrays
296        .iter()
297        .map(|x| x.as_run::<R>())
298        .filter(|x| !x.run_ends().is_empty())
299        .collect();
300
301    // The run ends need to be adjusted by the sum of the lengths of the previous arrays.
302    let needed_run_end_adjustments = std::iter::once(R::default_value())
303        .chain(
304            run_arrays
305                .iter()
306                .scan(R::default_value(), |acc, run_array| {
307                    *acc = *acc + *run_array.run_ends().values().last().unwrap();
308                    Some(*acc)
309                }),
310        )
311        .collect::<Vec<_>>();
312
313    // This works out nicely to be the total (logical) length of the resulting array.
314    let total_len = needed_run_end_adjustments.last().unwrap().as_usize();
315
316    let run_ends_array =
317        PrimitiveArray::<R>::from_iter_values(run_arrays.iter().enumerate().flat_map(
318            move |(i, run_array)| {
319                let adjustment = needed_run_end_adjustments[i];
320                run_array
321                    .run_ends()
322                    .values()
323                    .iter()
324                    .map(move |run_end| *run_end + adjustment)
325            },
326        ));
327
328    let all_values = concat(
329        &run_arrays
330            .iter()
331            .map(|x| x.values().as_ref())
332            .collect::<Vec<_>>(),
333    )?;
334
335    let builder = ArrayDataBuilder::new(run_arrays[0].data_type().clone())
336        .len(total_len)
337        .child_data(vec![run_ends_array.into_data(), all_values.into_data()]);
338
339    // `build_unchecked` is used to avoid recursive validation of child arrays.
340    let array_data = unsafe { builder.build_unchecked() };
341    array_data.validate_data()?;
342
343    Ok(Arc::<RunArray<R>>::new(array_data.into()))
344}
345
346macro_rules! dict_helper {
347    ($t:ty, $arrays:expr) => {
348        return Ok(Arc::new(concat_dictionaries::<$t>($arrays)?) as _)
349    };
350}
351
352macro_rules! primitive_concat {
353    ($t:ty, $arrays:expr) => {
354        return Ok(Arc::new(concat_primitives::<$t>($arrays)?) as _)
355    };
356}
357
358fn get_capacity(arrays: &[&dyn Array], data_type: &DataType) -> Capacities {
359    match data_type {
360        DataType::Utf8 => binary_capacity::<Utf8Type>(arrays),
361        DataType::LargeUtf8 => binary_capacity::<LargeUtf8Type>(arrays),
362        DataType::Binary => binary_capacity::<BinaryType>(arrays),
363        DataType::LargeBinary => binary_capacity::<LargeBinaryType>(arrays),
364        DataType::FixedSizeList(_, _) => fixed_size_list_capacity(arrays, data_type),
365        _ => Capacities::Array(arrays.iter().map(|a| a.len()).sum()),
366    }
367}
368
369/// Concatenate multiple [Array] of the same type into a single [ArrayRef].
370pub fn concat(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError> {
371    if arrays.is_empty() {
372        return Err(ArrowError::ComputeError(
373            "concat requires input of at least one array".to_string(),
374        ));
375    } else if arrays.len() == 1 {
376        let array = arrays[0];
377        return Ok(array.slice(0, array.len()));
378    }
379
380    let d = arrays[0].data_type();
381    if arrays.iter().skip(1).any(|array| array.data_type() != d) {
382        // Create error message with up to 10 unique data types in the order they appear
383        let error_message = {
384            // 10 max unique data types to print and another 1 to know if there are more
385            let mut unique_data_types = HashSet::with_capacity(11);
386
387            let mut error_message =
388                format!("It is not possible to concatenate arrays of different data types ({d}");
389            unique_data_types.insert(d);
390
391            for array in arrays {
392                let is_unique = unique_data_types.insert(array.data_type());
393
394                if unique_data_types.len() == 11 {
395                    error_message.push_str(", ...");
396                    break;
397                }
398
399                if is_unique {
400                    error_message.push_str(", ");
401                    error_message.push_str(&array.data_type().to_string());
402                }
403            }
404
405            error_message.push_str(").");
406
407            error_message
408        };
409
410        return Err(ArrowError::InvalidArgumentError(error_message));
411    }
412
413    downcast_primitive! {
414        d => (primitive_concat, arrays),
415        DataType::Boolean => concat_boolean(arrays),
416        DataType::Dictionary(k, _) => {
417            downcast_integer! {
418                k.as_ref() => (dict_helper, arrays),
419                _ => unreachable!("illegal dictionary key type {k}")
420            }
421        }
422        DataType::List(field) => concat_lists::<i32>(arrays, field),
423        DataType::LargeList(field) => concat_lists::<i64>(arrays, field),
424        DataType::Struct(fields) => concat_structs(arrays, fields),
425        DataType::Utf8 => concat_bytes::<Utf8Type>(arrays),
426        DataType::LargeUtf8 => concat_bytes::<LargeUtf8Type>(arrays),
427        DataType::Binary => concat_bytes::<BinaryType>(arrays),
428        DataType::LargeBinary => concat_bytes::<LargeBinaryType>(arrays),
429        DataType::RunEndEncoded(r, _) => {
430            // Handle RunEndEncoded arrays with special concat function
431            // We need to downcast based on the run end type
432            match r.data_type() {
433                DataType::Int16 => concat_run_arrays::<Int16Type>(arrays),
434                DataType::Int32 => concat_run_arrays::<Int32Type>(arrays),
435                DataType::Int64 => concat_run_arrays::<Int64Type>(arrays),
436                _ => unreachable!("Unsupported run end index type: {r:?}"),
437            }
438        }
439        DataType::Utf8View => concat_byte_view::<StringViewType>(arrays),
440        DataType::BinaryView => concat_byte_view::<BinaryViewType>(arrays),
441        _ => {
442            let capacity = get_capacity(arrays, d);
443            concat_fallback(arrays, capacity)
444        }
445    }
446}
447
448/// Concatenates arrays using MutableArrayData
449///
450/// This will naively concatenate dictionaries
451fn concat_fallback(arrays: &[&dyn Array], capacity: Capacities) -> Result<ArrayRef, ArrowError> {
452    let array_data: Vec<_> = arrays.iter().map(|a| a.to_data()).collect::<Vec<_>>();
453    let array_data = array_data.iter().collect();
454    let mut mutable = MutableArrayData::with_capacities(array_data, false, capacity);
455
456    for (i, a) in arrays.iter().enumerate() {
457        mutable.extend(i, 0, a.len())
458    }
459
460    Ok(make_array(mutable.freeze()))
461}
462
463/// Concatenates `batches` together into a single [`RecordBatch`].
464///
465/// The output batch has the specified `schemas`; The schema of the
466/// input are ignored.
467///
468/// Returns an error if the types of underlying arrays are different.
469pub fn concat_batches<'a>(
470    schema: &SchemaRef,
471    input_batches: impl IntoIterator<Item = &'a RecordBatch>,
472) -> Result<RecordBatch, ArrowError> {
473    // When schema is empty, sum the number of the rows of all batches
474    if schema.fields().is_empty() {
475        let num_rows: usize = input_batches.into_iter().map(RecordBatch::num_rows).sum();
476        let mut options = RecordBatchOptions::default();
477        options.row_count = Some(num_rows);
478        return RecordBatch::try_new_with_options(schema.clone(), vec![], &options);
479    }
480
481    let batches: Vec<&RecordBatch> = input_batches.into_iter().collect();
482    if batches.is_empty() {
483        return Ok(RecordBatch::new_empty(schema.clone()));
484    }
485    let field_num = schema.fields().len();
486    let mut arrays = Vec::with_capacity(field_num);
487    for i in 0..field_num {
488        let array = concat(
489            &batches
490                .iter()
491                .map(|batch| batch.column(i).as_ref())
492                .collect::<Vec<_>>(),
493        )?;
494        arrays.push(array);
495    }
496    RecordBatch::try_new(schema.clone(), arrays)
497}
498
499#[cfg(test)]
500mod tests {
501    use super::*;
502    use arrow_array::builder::{GenericListBuilder, StringDictionaryBuilder};
503    use arrow_schema::{Field, Schema};
504    use std::fmt::Debug;
505
506    #[test]
507    fn test_concat_empty_vec() {
508        let re = concat(&[]);
509        assert!(re.is_err());
510    }
511
512    #[test]
513    fn test_concat_batches_no_columns() {
514        // Test concat using empty schema / batches without columns
515        let schema = Arc::new(Schema::empty());
516
517        let mut options = RecordBatchOptions::default();
518        options.row_count = Some(100);
519        let batch = RecordBatch::try_new_with_options(schema.clone(), vec![], &options).unwrap();
520        // put in 2 batches of 100 rows each
521        let re = concat_batches(&schema, &[batch.clone(), batch]).unwrap();
522
523        assert_eq!(re.num_rows(), 200);
524    }
525
526    #[test]
527    fn test_concat_one_element_vec() {
528        let arr = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
529            Some(-1),
530            Some(2),
531            None,
532        ])) as ArrayRef;
533        let result = concat(&[arr.as_ref()]).unwrap();
534        assert_eq!(
535            &arr, &result,
536            "concatenating single element array gives back the same result"
537        );
538    }
539
540    #[test]
541    fn test_concat_incompatible_datatypes() {
542        let re = concat(&[
543            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(2), None]),
544            // 2 string to make sure we only mention unique types
545            &StringArray::from(vec![Some("hello"), Some("bar"), Some("world")]),
546            &StringArray::from(vec![Some("hey"), Some(""), Some("you")]),
547            // Another type to make sure we are showing all the incompatible types
548            &PrimitiveArray::<Int32Type>::from(vec![Some(-1), Some(2), None]),
549        ]);
550
551        assert_eq!(re.unwrap_err().to_string(), "Invalid argument error: It is not possible to concatenate arrays of different data types (Int64, Utf8, Int32).");
552    }
553
554    #[test]
555    fn test_concat_10_incompatible_datatypes_should_include_all_of_them() {
556        let re = concat(&[
557            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(2), None]),
558            // 2 string to make sure we only mention unique types
559            &StringArray::from(vec![Some("hello"), Some("bar"), Some("world")]),
560            &StringArray::from(vec![Some("hey"), Some(""), Some("you")]),
561            // Another type to make sure we are showing all the incompatible types
562            &PrimitiveArray::<Int32Type>::from(vec![Some(-1), Some(2), None]),
563            &PrimitiveArray::<Int8Type>::from(vec![Some(-1), Some(2), None]),
564            &PrimitiveArray::<Int16Type>::from(vec![Some(-1), Some(2), None]),
565            &PrimitiveArray::<UInt8Type>::from(vec![Some(1), Some(2), None]),
566            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
567            &PrimitiveArray::<UInt32Type>::from(vec![Some(1), Some(2), None]),
568            // Non unique
569            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
570            &PrimitiveArray::<UInt64Type>::from(vec![Some(1), Some(2), None]),
571            &PrimitiveArray::<Float32Type>::from(vec![Some(1.0), Some(2.0), None]),
572        ]);
573
574        assert_eq!(re.unwrap_err().to_string(), "Invalid argument error: It is not possible to concatenate arrays of different data types (Int64, Utf8, Int32, Int8, Int16, UInt8, UInt16, UInt32, UInt64, Float32).");
575    }
576
577    #[test]
578    fn test_concat_11_incompatible_datatypes_should_only_include_10() {
579        let re = concat(&[
580            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(2), None]),
581            // 2 string to make sure we only mention unique types
582            &StringArray::from(vec![Some("hello"), Some("bar"), Some("world")]),
583            &StringArray::from(vec![Some("hey"), Some(""), Some("you")]),
584            // Another type to make sure we are showing all the incompatible types
585            &PrimitiveArray::<Int32Type>::from(vec![Some(-1), Some(2), None]),
586            &PrimitiveArray::<Int8Type>::from(vec![Some(-1), Some(2), None]),
587            &PrimitiveArray::<Int16Type>::from(vec![Some(-1), Some(2), None]),
588            &PrimitiveArray::<UInt8Type>::from(vec![Some(1), Some(2), None]),
589            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
590            &PrimitiveArray::<UInt32Type>::from(vec![Some(1), Some(2), None]),
591            // Non unique
592            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
593            &PrimitiveArray::<UInt64Type>::from(vec![Some(1), Some(2), None]),
594            &PrimitiveArray::<Float32Type>::from(vec![Some(1.0), Some(2.0), None]),
595            &PrimitiveArray::<Float64Type>::from(vec![Some(1.0), Some(2.0), None]),
596        ]);
597
598        assert_eq!(re.unwrap_err().to_string(), "Invalid argument error: It is not possible to concatenate arrays of different data types (Int64, Utf8, Int32, Int8, Int16, UInt8, UInt16, UInt32, UInt64, Float32, ...).");
599    }
600
601    #[test]
602    fn test_concat_13_incompatible_datatypes_should_not_include_all_of_them() {
603        let re = concat(&[
604            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(2), None]),
605            // 2 string to make sure we only mention unique types
606            &StringArray::from(vec![Some("hello"), Some("bar"), Some("world")]),
607            &StringArray::from(vec![Some("hey"), Some(""), Some("you")]),
608            // Another type to make sure we are showing all the incompatible types
609            &PrimitiveArray::<Int32Type>::from(vec![Some(-1), Some(2), None]),
610            &PrimitiveArray::<Int8Type>::from(vec![Some(-1), Some(2), None]),
611            &PrimitiveArray::<Int16Type>::from(vec![Some(-1), Some(2), None]),
612            &PrimitiveArray::<UInt8Type>::from(vec![Some(1), Some(2), None]),
613            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
614            &PrimitiveArray::<UInt32Type>::from(vec![Some(1), Some(2), None]),
615            // Non unique
616            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
617            &PrimitiveArray::<UInt64Type>::from(vec![Some(1), Some(2), None]),
618            &PrimitiveArray::<Float32Type>::from(vec![Some(1.0), Some(2.0), None]),
619            &PrimitiveArray::<Float64Type>::from(vec![Some(1.0), Some(2.0), None]),
620            &PrimitiveArray::<Float16Type>::new_null(3),
621            &BooleanArray::from(vec![Some(true), Some(false), None]),
622        ]);
623
624        assert_eq!(re.unwrap_err().to_string(), "Invalid argument error: It is not possible to concatenate arrays of different data types (Int64, Utf8, Int32, Int8, Int16, UInt8, UInt16, UInt32, UInt64, Float32, ...).");
625    }
626
627    #[test]
628    fn test_concat_string_arrays() {
629        let arr = concat(&[
630            &StringArray::from(vec!["hello", "world"]),
631            &StringArray::from(vec!["2", "3", "4"]),
632            &StringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]),
633        ])
634        .unwrap();
635
636        let expected_output = Arc::new(StringArray::from(vec![
637            Some("hello"),
638            Some("world"),
639            Some("2"),
640            Some("3"),
641            Some("4"),
642            Some("foo"),
643            Some("bar"),
644            None,
645            Some("baz"),
646        ])) as ArrayRef;
647
648        assert_eq!(&arr, &expected_output);
649    }
650
651    #[test]
652    fn test_concat_string_view_arrays() {
653        let arr = concat(&[
654            &StringViewArray::from(vec!["helloxxxxxxxxxxa", "world____________"]),
655            &StringViewArray::from(vec!["helloxxxxxxxxxxy", "3", "4"]),
656            &StringViewArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]),
657        ])
658        .unwrap();
659
660        let expected_output = Arc::new(StringViewArray::from(vec![
661            Some("helloxxxxxxxxxxa"),
662            Some("world____________"),
663            Some("helloxxxxxxxxxxy"),
664            Some("3"),
665            Some("4"),
666            Some("foo"),
667            Some("bar"),
668            None,
669            Some("baz"),
670        ])) as ArrayRef;
671
672        assert_eq!(&arr, &expected_output);
673    }
674
675    #[test]
676    fn test_concat_primitive_arrays() {
677        let arr = concat(&[
678            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(-1), Some(2), None, None]),
679            &PrimitiveArray::<Int64Type>::from(vec![Some(101), Some(102), Some(103), None]),
680            &PrimitiveArray::<Int64Type>::from(vec![Some(256), Some(512), Some(1024)]),
681        ])
682        .unwrap();
683
684        let expected_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
685            Some(-1),
686            Some(-1),
687            Some(2),
688            None,
689            None,
690            Some(101),
691            Some(102),
692            Some(103),
693            None,
694            Some(256),
695            Some(512),
696            Some(1024),
697        ])) as ArrayRef;
698
699        assert_eq!(&arr, &expected_output);
700    }
701
702    #[test]
703    fn test_concat_primitive_array_slices() {
704        let input_1 =
705            PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(-1), Some(2), None, None])
706                .slice(1, 3);
707
708        let input_2 =
709            PrimitiveArray::<Int64Type>::from(vec![Some(101), Some(102), Some(103), None])
710                .slice(1, 3);
711        let arr = concat(&[&input_1, &input_2]).unwrap();
712
713        let expected_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
714            Some(-1),
715            Some(2),
716            None,
717            Some(102),
718            Some(103),
719            None,
720        ])) as ArrayRef;
721
722        assert_eq!(&arr, &expected_output);
723    }
724
725    #[test]
726    fn test_concat_boolean_primitive_arrays() {
727        let arr = concat(&[
728            &BooleanArray::from(vec![
729                Some(true),
730                Some(true),
731                Some(false),
732                None,
733                None,
734                Some(false),
735            ]),
736            &BooleanArray::from(vec![None, Some(false), Some(true), Some(false)]),
737        ])
738        .unwrap();
739
740        let expected_output = Arc::new(BooleanArray::from(vec![
741            Some(true),
742            Some(true),
743            Some(false),
744            None,
745            None,
746            Some(false),
747            None,
748            Some(false),
749            Some(true),
750            Some(false),
751        ])) as ArrayRef;
752
753        assert_eq!(&arr, &expected_output);
754    }
755
756    #[test]
757    fn test_concat_primitive_list_arrays() {
758        let list1 = vec![
759            Some(vec![Some(-1), Some(-1), Some(2), None, None]),
760            Some(vec![]),
761            None,
762            Some(vec![Some(10)]),
763        ];
764        let list1_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone());
765
766        let list2 = vec![
767            None,
768            Some(vec![Some(100), None, Some(101)]),
769            Some(vec![Some(102)]),
770        ];
771        let list2_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone());
772
773        let list3 = vec![Some(vec![Some(1000), Some(1001)])];
774        let list3_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list3.clone());
775
776        let array_result = concat(&[&list1_array, &list2_array, &list3_array]).unwrap();
777
778        let expected = list1.into_iter().chain(list2).chain(list3);
779        let array_expected = ListArray::from_iter_primitive::<Int64Type, _, _>(expected);
780
781        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
782    }
783
784    #[test]
785    fn test_concat_primitive_list_arrays_slices() {
786        let list1 = vec![
787            Some(vec![Some(-1), Some(-1), Some(2), None, None]),
788            Some(vec![]), // In slice
789            None,         // In slice
790            Some(vec![Some(10)]),
791        ];
792        let list1_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone());
793        let list1_array = list1_array.slice(1, 2);
794        let list1_values = list1.into_iter().skip(1).take(2);
795
796        let list2 = vec![
797            None,
798            Some(vec![Some(100), None, Some(101)]),
799            Some(vec![Some(102)]),
800        ];
801        let list2_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone());
802
803        // verify that this test covers the case when the first offset is non zero
804        assert!(list1_array.offsets()[0].as_usize() > 0);
805        let array_result = concat(&[&list1_array, &list2_array]).unwrap();
806
807        let expected = list1_values.chain(list2);
808        let array_expected = ListArray::from_iter_primitive::<Int64Type, _, _>(expected);
809
810        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
811    }
812
813    #[test]
814    fn test_concat_primitive_list_arrays_sliced_lengths() {
815        let list1 = vec![
816            Some(vec![Some(-1), Some(-1), Some(2), None, None]), // In slice
817            Some(vec![]),                                        // In slice
818            None,                                                // In slice
819            Some(vec![Some(10)]),
820        ];
821        let list1_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone());
822        let list1_array = list1_array.slice(0, 3); // no offset, but not all values
823        let list1_values = list1.into_iter().take(3);
824
825        let list2 = vec![
826            None,
827            Some(vec![Some(100), None, Some(101)]),
828            Some(vec![Some(102)]),
829        ];
830        let list2_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone());
831
832        // verify that this test covers the case when the first offset is zero, but the
833        // last offset doesn't cover the entire array
834        assert_eq!(list1_array.offsets()[0].as_usize(), 0);
835        assert!(list1_array.offsets().last().unwrap().as_usize() < list1_array.values().len());
836        let array_result = concat(&[&list1_array, &list2_array]).unwrap();
837
838        let expected = list1_values.chain(list2);
839        let array_expected = ListArray::from_iter_primitive::<Int64Type, _, _>(expected);
840
841        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
842    }
843
844    #[test]
845    fn test_concat_primitive_fixed_size_list_arrays() {
846        let list1 = vec![
847            Some(vec![Some(-1), None]),
848            None,
849            Some(vec![Some(10), Some(20)]),
850        ];
851        let list1_array =
852            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone(), 2);
853
854        let list2 = vec![
855            None,
856            Some(vec![Some(100), None]),
857            Some(vec![Some(102), Some(103)]),
858        ];
859        let list2_array =
860            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone(), 2);
861
862        let list3 = vec![Some(vec![Some(1000), Some(1001)])];
863        let list3_array =
864            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(list3.clone(), 2);
865
866        let array_result = concat(&[&list1_array, &list2_array, &list3_array]).unwrap();
867
868        let expected = list1.into_iter().chain(list2).chain(list3);
869        let array_expected =
870            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(expected, 2);
871
872        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
873    }
874
875    #[test]
876    fn test_concat_struct_arrays() {
877        let field = Arc::new(Field::new("field", DataType::Int64, true));
878        let input_primitive_1: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
879            Some(-1),
880            Some(-1),
881            Some(2),
882            None,
883            None,
884        ]));
885        let input_struct_1 = StructArray::from(vec![(field.clone(), input_primitive_1)]);
886
887        let input_primitive_2: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
888            Some(101),
889            Some(102),
890            Some(103),
891            None,
892        ]));
893        let input_struct_2 = StructArray::from(vec![(field.clone(), input_primitive_2)]);
894
895        let input_primitive_3: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
896            Some(256),
897            Some(512),
898            Some(1024),
899        ]));
900        let input_struct_3 = StructArray::from(vec![(field, input_primitive_3)]);
901
902        let arr = concat(&[&input_struct_1, &input_struct_2, &input_struct_3]).unwrap();
903
904        let expected_primitive_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
905            Some(-1),
906            Some(-1),
907            Some(2),
908            None,
909            None,
910            Some(101),
911            Some(102),
912            Some(103),
913            None,
914            Some(256),
915            Some(512),
916            Some(1024),
917        ])) as ArrayRef;
918
919        let actual_primitive = arr
920            .as_any()
921            .downcast_ref::<StructArray>()
922            .unwrap()
923            .column(0);
924        assert_eq!(actual_primitive, &expected_primitive_output);
925    }
926
927    #[test]
928    fn test_concat_struct_array_slices() {
929        let field = Arc::new(Field::new("field", DataType::Int64, true));
930        let input_primitive_1: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
931            Some(-1),
932            Some(-1),
933            Some(2),
934            None,
935            None,
936        ]));
937        let input_struct_1 = StructArray::from(vec![(field.clone(), input_primitive_1)]);
938
939        let input_primitive_2: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
940            Some(101),
941            Some(102),
942            Some(103),
943            None,
944        ]));
945        let input_struct_2 = StructArray::from(vec![(field, input_primitive_2)]);
946
947        let arr = concat(&[&input_struct_1.slice(1, 3), &input_struct_2.slice(1, 2)]).unwrap();
948
949        let expected_primitive_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
950            Some(-1),
951            Some(2),
952            None,
953            Some(102),
954            Some(103),
955        ])) as ArrayRef;
956
957        let actual_primitive = arr
958            .as_any()
959            .downcast_ref::<StructArray>()
960            .unwrap()
961            .column(0);
962        assert_eq!(actual_primitive, &expected_primitive_output);
963    }
964
965    #[test]
966    fn test_concat_struct_arrays_no_nulls() {
967        let input_1a = vec![1, 2, 3];
968        let input_1b = vec!["one", "two", "three"];
969        let input_2a = vec![4, 5, 6, 7];
970        let input_2b = vec!["four", "five", "six", "seven"];
971
972        let struct_from_primitives = |ints: Vec<i64>, strings: Vec<&str>| {
973            StructArray::try_from(vec![
974                ("ints", Arc::new(Int64Array::from(ints)) as _),
975                ("strings", Arc::new(StringArray::from(strings)) as _),
976            ])
977        };
978
979        let expected_output = struct_from_primitives(
980            [input_1a.clone(), input_2a.clone()].concat(),
981            [input_1b.clone(), input_2b.clone()].concat(),
982        )
983        .unwrap();
984
985        let input_1 = struct_from_primitives(input_1a, input_1b).unwrap();
986        let input_2 = struct_from_primitives(input_2a, input_2b).unwrap();
987
988        let arr = concat(&[&input_1, &input_2]).unwrap();
989        let struct_result = arr.as_struct();
990
991        assert_eq!(struct_result, &expected_output);
992        assert_eq!(arr.null_count(), 0);
993    }
994
995    #[test]
996    fn test_string_array_slices() {
997        let input_1 = StringArray::from(vec!["hello", "A", "B", "C"]);
998        let input_2 = StringArray::from(vec!["world", "D", "E", "Z"]);
999
1000        let arr = concat(&[&input_1.slice(1, 3), &input_2.slice(1, 2)]).unwrap();
1001
1002        let expected_output = StringArray::from(vec!["A", "B", "C", "D", "E"]);
1003
1004        let actual_output = arr.as_any().downcast_ref::<StringArray>().unwrap();
1005        assert_eq!(actual_output, &expected_output);
1006    }
1007
1008    #[test]
1009    fn test_string_array_with_null_slices() {
1010        let input_1 = StringArray::from(vec![Some("hello"), None, Some("A"), Some("C")]);
1011        let input_2 = StringArray::from(vec![None, Some("world"), Some("D"), None]);
1012
1013        let arr = concat(&[&input_1.slice(1, 3), &input_2.slice(1, 2)]).unwrap();
1014
1015        let expected_output =
1016            StringArray::from(vec![None, Some("A"), Some("C"), Some("world"), Some("D")]);
1017
1018        let actual_output = arr.as_any().downcast_ref::<StringArray>().unwrap();
1019        assert_eq!(actual_output, &expected_output);
1020    }
1021
1022    fn collect_string_dictionary(array: &DictionaryArray<Int32Type>) -> Vec<Option<&str>> {
1023        let concrete = array.downcast_dict::<StringArray>().unwrap();
1024        concrete.into_iter().collect()
1025    }
1026
1027    #[test]
1028    fn test_string_dictionary_array() {
1029        let input_1: DictionaryArray<Int32Type> = vec!["hello", "A", "B", "hello", "hello", "C"]
1030            .into_iter()
1031            .collect();
1032        let input_2: DictionaryArray<Int32Type> = vec!["hello", "E", "E", "hello", "F", "E"]
1033            .into_iter()
1034            .collect();
1035
1036        let expected: Vec<_> = vec![
1037            "hello", "A", "B", "hello", "hello", "C", "hello", "E", "E", "hello", "F", "E",
1038        ]
1039        .into_iter()
1040        .map(Some)
1041        .collect();
1042
1043        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1044        let dictionary = concat.as_dictionary::<Int32Type>();
1045        let actual = collect_string_dictionary(dictionary);
1046        assert_eq!(actual, expected);
1047
1048        // Should have concatenated inputs together
1049        assert_eq!(
1050            dictionary.values().len(),
1051            input_1.values().len() + input_2.values().len(),
1052        )
1053    }
1054
1055    #[test]
1056    fn test_string_dictionary_array_nulls() {
1057        let input_1: DictionaryArray<Int32Type> = vec![Some("foo"), Some("bar"), None, Some("fiz")]
1058            .into_iter()
1059            .collect();
1060        let input_2: DictionaryArray<Int32Type> = vec![None].into_iter().collect();
1061        let expected = vec![Some("foo"), Some("bar"), None, Some("fiz"), None];
1062
1063        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1064        let dictionary = concat.as_dictionary::<Int32Type>();
1065        let actual = collect_string_dictionary(dictionary);
1066        assert_eq!(actual, expected);
1067
1068        // Should have concatenated inputs together
1069        assert_eq!(
1070            dictionary.values().len(),
1071            input_1.values().len() + input_2.values().len(),
1072        )
1073    }
1074
1075    #[test]
1076    fn test_string_dictionary_array_nulls_in_values() {
1077        let input_1_keys = Int32Array::from_iter_values([0, 2, 1, 3]);
1078        let input_1_values = StringArray::from(vec![Some("foo"), None, Some("bar"), Some("fiz")]);
1079        let input_1 = DictionaryArray::new(input_1_keys, Arc::new(input_1_values));
1080
1081        let input_2_keys = Int32Array::from_iter_values([0]);
1082        let input_2_values = StringArray::from(vec![None, Some("hello")]);
1083        let input_2 = DictionaryArray::new(input_2_keys, Arc::new(input_2_values));
1084
1085        let expected = vec![Some("foo"), Some("bar"), None, Some("fiz"), None];
1086
1087        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1088        let dictionary = concat.as_dictionary::<Int32Type>();
1089        let actual = collect_string_dictionary(dictionary);
1090        assert_eq!(actual, expected);
1091    }
1092
1093    #[test]
1094    fn test_string_dictionary_merge() {
1095        let mut builder = StringDictionaryBuilder::<Int32Type>::new();
1096        for i in 0..20 {
1097            builder.append(i.to_string()).unwrap();
1098        }
1099        let input_1 = builder.finish();
1100
1101        let mut builder = StringDictionaryBuilder::<Int32Type>::new();
1102        for i in 0..30 {
1103            builder.append(i.to_string()).unwrap();
1104        }
1105        let input_2 = builder.finish();
1106
1107        let expected: Vec<_> = (0..20).chain(0..30).map(|x| x.to_string()).collect();
1108        let expected: Vec<_> = expected.iter().map(|x| Some(x.as_str())).collect();
1109
1110        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1111        let dictionary = concat.as_dictionary::<Int32Type>();
1112        let actual = collect_string_dictionary(dictionary);
1113        assert_eq!(actual, expected);
1114
1115        // Should have merged inputs together
1116        // Not 30 as this is done on a best-effort basis
1117        let values_len = dictionary.values().len();
1118        assert!((30..40).contains(&values_len), "{values_len}")
1119    }
1120
1121    #[test]
1122    fn test_primitive_dictionary_merge() {
1123        // Same value repeated 5 times.
1124        let keys = vec![1; 5];
1125        let values = (10..20).collect::<Vec<_>>();
1126        let dict = DictionaryArray::new(
1127            Int8Array::from(keys.clone()),
1128            Arc::new(Int32Array::from(values.clone())),
1129        );
1130        let other = DictionaryArray::new(
1131            Int8Array::from(keys.clone()),
1132            Arc::new(Int32Array::from(values.clone())),
1133        );
1134
1135        let result_same_dictionary = concat(&[&dict, &dict]).unwrap();
1136        // Verify pointer equality check succeeds, and therefore the
1137        // dictionaries are not merged. A single values buffer should be reused
1138        // in this case.
1139        assert!(dict.values().to_data().ptr_eq(
1140            &result_same_dictionary
1141                .as_dictionary::<Int8Type>()
1142                .values()
1143                .to_data()
1144        ));
1145        assert_eq!(
1146            result_same_dictionary
1147                .as_dictionary::<Int8Type>()
1148                .values()
1149                .len(),
1150            values.len(),
1151        );
1152
1153        let result_cloned_dictionary = concat(&[&dict, &other]).unwrap();
1154        // Should have only 1 underlying value since all keys reference it.
1155        assert_eq!(
1156            result_cloned_dictionary
1157                .as_dictionary::<Int8Type>()
1158                .values()
1159                .len(),
1160            1
1161        );
1162    }
1163
1164    #[test]
1165    fn test_concat_string_sizes() {
1166        let a: LargeStringArray = ((0..150).map(|_| Some("foo"))).collect();
1167        let b: LargeStringArray = ((0..150).map(|_| Some("foo"))).collect();
1168        let c = LargeStringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]);
1169        // 150 * 3 = 450
1170        // 150 * 3 = 450
1171        // 3 * 3   = 9
1172        // ------------+
1173        // 909
1174        // closest 64 byte aligned cap = 960
1175
1176        let arr = concat(&[&a, &b, &c]).unwrap();
1177        // this would have been 1280 if we did not precompute the value lengths.
1178        assert_eq!(arr.to_data().buffers()[1].capacity(), 960);
1179    }
1180
1181    #[test]
1182    fn test_dictionary_concat_reuse() {
1183        let array: DictionaryArray<Int8Type> = vec!["a", "a", "b", "c"].into_iter().collect();
1184        let copy: DictionaryArray<Int8Type> = array.clone();
1185
1186        // dictionary is "a", "b", "c"
1187        assert_eq!(
1188            array.values(),
1189            &(Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef)
1190        );
1191        assert_eq!(array.keys(), &Int8Array::from(vec![0, 0, 1, 2]));
1192
1193        // concatenate it with itself
1194        let combined = concat(&[&copy as _, &array as _]).unwrap();
1195        let combined = combined.as_dictionary::<Int8Type>();
1196
1197        assert_eq!(
1198            combined.values(),
1199            &(Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef),
1200            "Actual: {combined:#?}"
1201        );
1202
1203        assert_eq!(
1204            combined.keys(),
1205            &Int8Array::from(vec![0, 0, 1, 2, 0, 0, 1, 2])
1206        );
1207
1208        // Should have reused the dictionary
1209        assert!(array
1210            .values()
1211            .to_data()
1212            .ptr_eq(&combined.values().to_data()));
1213        assert!(copy.values().to_data().ptr_eq(&combined.values().to_data()));
1214
1215        let new: DictionaryArray<Int8Type> = vec!["d"].into_iter().collect();
1216        let combined = concat(&[&copy as _, &array as _, &new as _]).unwrap();
1217        let com = combined.as_dictionary::<Int8Type>();
1218
1219        // Should not have reused the dictionary
1220        assert!(!array.values().to_data().ptr_eq(&com.values().to_data()));
1221        assert!(!copy.values().to_data().ptr_eq(&com.values().to_data()));
1222        assert!(!new.values().to_data().ptr_eq(&com.values().to_data()));
1223    }
1224
1225    #[test]
1226    fn concat_record_batches() {
1227        let schema = Arc::new(Schema::new(vec![
1228            Field::new("a", DataType::Int32, false),
1229            Field::new("b", DataType::Utf8, false),
1230        ]));
1231        let batch1 = RecordBatch::try_new(
1232            schema.clone(),
1233            vec![
1234                Arc::new(Int32Array::from(vec![1, 2])),
1235                Arc::new(StringArray::from(vec!["a", "b"])),
1236            ],
1237        )
1238        .unwrap();
1239        let batch2 = RecordBatch::try_new(
1240            schema.clone(),
1241            vec![
1242                Arc::new(Int32Array::from(vec![3, 4])),
1243                Arc::new(StringArray::from(vec!["c", "d"])),
1244            ],
1245        )
1246        .unwrap();
1247        let new_batch = concat_batches(&schema, [&batch1, &batch2]).unwrap();
1248        assert_eq!(new_batch.schema().as_ref(), schema.as_ref());
1249        assert_eq!(2, new_batch.num_columns());
1250        assert_eq!(4, new_batch.num_rows());
1251        let new_batch_owned = concat_batches(&schema, &[batch1, batch2]).unwrap();
1252        assert_eq!(new_batch_owned.schema().as_ref(), schema.as_ref());
1253        assert_eq!(2, new_batch_owned.num_columns());
1254        assert_eq!(4, new_batch_owned.num_rows());
1255    }
1256
1257    #[test]
1258    fn concat_empty_record_batch() {
1259        let schema = Arc::new(Schema::new(vec![
1260            Field::new("a", DataType::Int32, false),
1261            Field::new("b", DataType::Utf8, false),
1262        ]));
1263        let batch = concat_batches(&schema, []).unwrap();
1264        assert_eq!(batch.schema().as_ref(), schema.as_ref());
1265        assert_eq!(0, batch.num_rows());
1266    }
1267
1268    #[test]
1269    fn concat_record_batches_of_different_schemas_but_compatible_data() {
1270        let schema1 = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
1271        // column names differ
1272        let schema2 = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, false)]));
1273        let batch1 = RecordBatch::try_new(
1274            schema1.clone(),
1275            vec![Arc::new(Int32Array::from(vec![1, 2]))],
1276        )
1277        .unwrap();
1278        let batch2 =
1279            RecordBatch::try_new(schema2, vec![Arc::new(Int32Array::from(vec![3, 4]))]).unwrap();
1280        // concat_batches simply uses the schema provided
1281        let batch = concat_batches(&schema1, [&batch1, &batch2]).unwrap();
1282        assert_eq!(batch.schema().as_ref(), schema1.as_ref());
1283        assert_eq!(4, batch.num_rows());
1284    }
1285
1286    #[test]
1287    fn concat_record_batches_of_different_schemas_incompatible_data() {
1288        let schema1 = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
1289        // column names differ
1290        let schema2 = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, false)]));
1291        let batch1 = RecordBatch::try_new(
1292            schema1.clone(),
1293            vec![Arc::new(Int32Array::from(vec![1, 2]))],
1294        )
1295        .unwrap();
1296        let batch2 = RecordBatch::try_new(
1297            schema2,
1298            vec![Arc::new(StringArray::from(vec!["foo", "bar"]))],
1299        )
1300        .unwrap();
1301
1302        let error = concat_batches(&schema1, [&batch1, &batch2]).unwrap_err();
1303        assert_eq!(error.to_string(), "Invalid argument error: It is not possible to concatenate arrays of different data types (Int32, Utf8).");
1304    }
1305
1306    #[test]
1307    fn concat_capacity() {
1308        let a = Int32Array::from_iter_values(0..100);
1309        let b = Int32Array::from_iter_values(10..20);
1310        let a = concat(&[&a, &b]).unwrap();
1311        let data = a.to_data();
1312        assert_eq!(data.buffers()[0].len(), 440);
1313        assert_eq!(data.buffers()[0].capacity(), 448); // Nearest multiple of 64
1314
1315        let a = concat(&[&a.slice(10, 20), &b]).unwrap();
1316        let data = a.to_data();
1317        assert_eq!(data.buffers()[0].len(), 120);
1318        assert_eq!(data.buffers()[0].capacity(), 128); // Nearest multiple of 64
1319
1320        let a = StringArray::from_iter_values(std::iter::repeat("foo").take(100));
1321        let b = StringArray::from(vec!["bingo", "bongo", "lorem", ""]);
1322
1323        let a = concat(&[&a, &b]).unwrap();
1324        let data = a.to_data();
1325        // (100 + 4 + 1) * size_of<i32>()
1326        assert_eq!(data.buffers()[0].len(), 420);
1327        assert_eq!(data.buffers()[0].capacity(), 448); // Nearest multiple of 64
1328
1329        // len("foo") * 100 + len("bingo") + len("bongo") + len("lorem")
1330        assert_eq!(data.buffers()[1].len(), 315);
1331        assert_eq!(data.buffers()[1].capacity(), 320); // Nearest multiple of 64
1332
1333        let a = concat(&[&a.slice(10, 40), &b]).unwrap();
1334        let data = a.to_data();
1335        // (40 + 4 + 5) * size_of<i32>()
1336        assert_eq!(data.buffers()[0].len(), 180);
1337        assert_eq!(data.buffers()[0].capacity(), 192); // Nearest multiple of 64
1338
1339        // len("foo") * 40 + len("bingo") + len("bongo") + len("lorem")
1340        assert_eq!(data.buffers()[1].len(), 135);
1341        assert_eq!(data.buffers()[1].capacity(), 192); // Nearest multiple of 64
1342
1343        let a = LargeBinaryArray::from_iter_values(std::iter::repeat(b"foo").take(100));
1344        let b = LargeBinaryArray::from_iter_values(std::iter::repeat(b"cupcakes").take(10));
1345
1346        let a = concat(&[&a, &b]).unwrap();
1347        let data = a.to_data();
1348        // (100 + 10 + 1) * size_of<i64>()
1349        assert_eq!(data.buffers()[0].len(), 888);
1350        assert_eq!(data.buffers()[0].capacity(), 896); // Nearest multiple of 64
1351
1352        // len("foo") * 100 + len("cupcakes") * 10
1353        assert_eq!(data.buffers()[1].len(), 380);
1354        assert_eq!(data.buffers()[1].capacity(), 384); // Nearest multiple of 64
1355
1356        let a = concat(&[&a.slice(10, 40), &b]).unwrap();
1357        let data = a.to_data();
1358        // (40 + 10 + 1) * size_of<i64>()
1359        assert_eq!(data.buffers()[0].len(), 408);
1360        assert_eq!(data.buffers()[0].capacity(), 448); // Nearest multiple of 64
1361
1362        // len("foo") * 40 + len("cupcakes") * 10
1363        assert_eq!(data.buffers()[1].len(), 200);
1364        assert_eq!(data.buffers()[1].capacity(), 256); // Nearest multiple of 64
1365    }
1366
1367    #[test]
1368    fn concat_sparse_nulls() {
1369        let values = StringArray::from_iter_values((0..100).map(|x| x.to_string()));
1370        let keys = Int32Array::from(vec![1; 10]);
1371        let dict_a = DictionaryArray::new(keys, Arc::new(values));
1372        let values = StringArray::new_null(0);
1373        let keys = Int32Array::new_null(10);
1374        let dict_b = DictionaryArray::new(keys, Arc::new(values));
1375        let array = concat(&[&dict_a, &dict_b]).unwrap();
1376        assert_eq!(array.null_count(), 10);
1377        assert_eq!(array.logical_null_count(), 10);
1378    }
1379
1380    #[test]
1381    fn concat_dictionary_list_array_simple() {
1382        let scalars = vec![
1383            create_single_row_list_of_dict(vec![Some("a")]),
1384            create_single_row_list_of_dict(vec![Some("a")]),
1385            create_single_row_list_of_dict(vec![Some("b")]),
1386        ];
1387
1388        let arrays = scalars
1389            .iter()
1390            .map(|a| a as &(dyn Array))
1391            .collect::<Vec<_>>();
1392        let concat_res = concat(arrays.as_slice()).unwrap();
1393
1394        let expected_list = create_list_of_dict(vec![
1395            // Row 1
1396            Some(vec![Some("a")]),
1397            Some(vec![Some("a")]),
1398            Some(vec![Some("b")]),
1399        ]);
1400
1401        let list = concat_res.as_list::<i32>();
1402
1403        // Assert that the list is equal to the expected list
1404        list.iter().zip(expected_list.iter()).for_each(|(a, b)| {
1405            assert_eq!(a, b);
1406        });
1407
1408        assert_dictionary_has_unique_values::<_, StringArray>(
1409            list.values().as_dictionary::<Int32Type>(),
1410        );
1411    }
1412
1413    #[test]
1414    fn concat_many_dictionary_list_arrays() {
1415        let number_of_unique_values = 8;
1416        let scalars = (0..80000)
1417            .map(|i| {
1418                create_single_row_list_of_dict(vec![Some(
1419                    (i % number_of_unique_values).to_string(),
1420                )])
1421            })
1422            .collect::<Vec<_>>();
1423
1424        let arrays = scalars
1425            .iter()
1426            .map(|a| a as &(dyn Array))
1427            .collect::<Vec<_>>();
1428        let concat_res = concat(arrays.as_slice()).unwrap();
1429
1430        let expected_list = create_list_of_dict(
1431            (0..80000)
1432                .map(|i| Some(vec![Some((i % number_of_unique_values).to_string())]))
1433                .collect::<Vec<_>>(),
1434        );
1435
1436        let list = concat_res.as_list::<i32>();
1437
1438        // Assert that the list is equal to the expected list
1439        list.iter().zip(expected_list.iter()).for_each(|(a, b)| {
1440            assert_eq!(a, b);
1441        });
1442
1443        assert_dictionary_has_unique_values::<_, StringArray>(
1444            list.values().as_dictionary::<Int32Type>(),
1445        );
1446    }
1447
1448    fn create_single_row_list_of_dict(
1449        list_items: Vec<Option<impl AsRef<str>>>,
1450    ) -> GenericListArray<i32> {
1451        let rows = list_items.into_iter().map(Some).collect();
1452
1453        create_list_of_dict(vec![rows])
1454    }
1455
1456    fn create_list_of_dict(
1457        rows: Vec<Option<Vec<Option<impl AsRef<str>>>>>,
1458    ) -> GenericListArray<i32> {
1459        let mut builder =
1460            GenericListBuilder::<i32, _>::new(StringDictionaryBuilder::<Int32Type>::new());
1461
1462        for row in rows {
1463            builder.append_option(row);
1464        }
1465
1466        builder.finish()
1467    }
1468
1469    fn assert_dictionary_has_unique_values<'a, K, V>(array: &'a DictionaryArray<K>)
1470    where
1471        K: ArrowDictionaryKeyType,
1472        V: Sync + Send + 'static,
1473        &'a V: ArrayAccessor + IntoIterator,
1474
1475        <&'a V as ArrayAccessor>::Item: Default + Clone + PartialEq + Debug + Ord,
1476        <&'a V as IntoIterator>::Item: Clone + PartialEq + Debug + Ord,
1477    {
1478        let dict = array.downcast_dict::<V>().unwrap();
1479        let mut values = dict.values().into_iter().collect::<Vec<_>>();
1480
1481        // remove duplicates must be sorted first so we can compare
1482        values.sort();
1483
1484        let mut unique_values = values.clone();
1485
1486        unique_values.dedup();
1487
1488        assert_eq!(
1489            values, unique_values,
1490            "There are duplicates in the value list (the value list here is sorted which is only for the assertion)"
1491        );
1492    }
1493
1494    // Test the simple case of concatenating two RunArrays
1495    #[test]
1496    fn test_concat_run_array() {
1497        // Create simple run arrays
1498        let run_ends1 = Int32Array::from(vec![2, 4]);
1499        let values1 = Int32Array::from(vec![10, 20]);
1500        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1501
1502        let run_ends2 = Int32Array::from(vec![1, 4]);
1503        let values2 = Int32Array::from(vec![30, 40]);
1504        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
1505
1506        // Concatenate the arrays - this should now work properly
1507        let result = concat(&[&array1, &array2]).unwrap();
1508        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1509
1510        // Check that the result has the correct length
1511        assert_eq!(result_run_array.len(), 8); // 4 + 4
1512
1513        // Check the run ends
1514        let run_ends = result_run_array.run_ends().values();
1515        assert_eq!(run_ends.len(), 4);
1516        assert_eq!(&[2, 4, 5, 8], run_ends);
1517
1518        // Check the values
1519        let values = result_run_array
1520            .values()
1521            .as_any()
1522            .downcast_ref::<Int32Array>()
1523            .unwrap();
1524        assert_eq!(values.len(), 4);
1525        assert_eq!(&[10, 20, 30, 40], values.values());
1526    }
1527
1528    #[test]
1529    fn test_concat_run_array_matching_first_last_value() {
1530        // Create a run array with run ends [2, 4, 7] and values [10, 20, 30]
1531        let run_ends1 = Int32Array::from(vec![2, 4, 7]);
1532        let values1 = Int32Array::from(vec![10, 20, 30]);
1533        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1534
1535        // Create another run array with run ends [3, 5] and values [30, 40]
1536        let run_ends2 = Int32Array::from(vec![3, 5]);
1537        let values2 = Int32Array::from(vec![30, 40]);
1538        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
1539
1540        // Concatenate the two arrays
1541        let result = concat(&[&array1, &array2]).unwrap();
1542        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1543
1544        // The result should have length 12 (7 + 5)
1545        assert_eq!(result_run_array.len(), 12);
1546
1547        // Check that the run ends are correct
1548        let run_ends = result_run_array.run_ends().values();
1549        assert_eq!(&[2, 4, 7, 10, 12], run_ends);
1550
1551        // Check that the values are correct
1552        assert_eq!(
1553            &[10, 20, 30, 30, 40],
1554            result_run_array
1555                .values()
1556                .as_any()
1557                .downcast_ref::<Int32Array>()
1558                .unwrap()
1559                .values()
1560        );
1561    }
1562
1563    #[test]
1564    fn test_concat_run_array_with_nulls() {
1565        // Create values array with nulls
1566        let values1 = Int32Array::from(vec![Some(10), None, Some(30)]);
1567        let run_ends1 = Int32Array::from(vec![2, 4, 7]);
1568        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1569
1570        // Create another run array with run ends [3, 5] and values [30, null]
1571        let values2 = Int32Array::from(vec![Some(30), None]);
1572        let run_ends2 = Int32Array::from(vec![3, 5]);
1573        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
1574
1575        // Concatenate the two arrays
1576        let result = concat(&[&array1, &array2]).unwrap();
1577        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1578
1579        // The result should have length 12 (7 + 5)
1580        assert_eq!(result_run_array.len(), 12);
1581
1582        // Get a reference to the run array itself for testing
1583
1584        // Just test the length and run ends without asserting specific values
1585        // This ensures the test passes while we work on full support for RunArray nulls
1586        assert_eq!(result_run_array.len(), 12); // 7 + 5
1587
1588        // Check that the run ends are correct
1589        let run_ends_values = result_run_array.run_ends().values();
1590        assert_eq!(&[2, 4, 7, 10, 12], run_ends_values);
1591
1592        // Check that the values are correct
1593        let expected = Int32Array::from(vec![Some(10), None, Some(30), Some(30), None]);
1594        let actual = result_run_array
1595            .values()
1596            .as_any()
1597            .downcast_ref::<Int32Array>()
1598            .unwrap();
1599        assert_eq!(actual.len(), expected.len());
1600        assert_eq!(actual.null_count(), expected.null_count());
1601        assert_eq!(actual.values(), expected.values());
1602    }
1603
1604    #[test]
1605    fn test_concat_run_array_single() {
1606        // Create a run array with run ends [2, 4] and values [10, 20]
1607        let run_ends1 = Int32Array::from(vec![2, 4]);
1608        let values1 = Int32Array::from(vec![10, 20]);
1609        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1610
1611        // Concatenate the single array
1612        let result = concat(&[&array1]).unwrap();
1613        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1614
1615        // The result should have length 4
1616        assert_eq!(result_run_array.len(), 4);
1617
1618        // Check that the run ends are correct
1619        let run_ends = result_run_array.run_ends().values();
1620        assert_eq!(&[2, 4], run_ends);
1621
1622        // Check that the values are correct
1623        assert_eq!(
1624            &[10, 20],
1625            result_run_array
1626                .values()
1627                .as_any()
1628                .downcast_ref::<Int32Array>()
1629                .unwrap()
1630                .values()
1631        );
1632    }
1633
1634    #[test]
1635    fn test_concat_run_array_with_3_arrays() {
1636        let run_ends1 = Int32Array::from(vec![2, 4]);
1637        let values1 = Int32Array::from(vec![10, 20]);
1638        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1639        let run_ends2 = Int32Array::from(vec![1, 4]);
1640        let values2 = Int32Array::from(vec![30, 40]);
1641        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
1642        let run_ends3 = Int32Array::from(vec![1, 4]);
1643        let values3 = Int32Array::from(vec![50, 60]);
1644        let array3 = RunArray::try_new(&run_ends3, &values3).unwrap();
1645
1646        // Concatenate the arrays
1647        let result = concat(&[&array1, &array2, &array3]).unwrap();
1648        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1649
1650        // Check that the result has the correct length
1651        assert_eq!(result_run_array.len(), 12); // 4 + 4 + 4
1652
1653        // Check the run ends
1654        let run_ends = result_run_array.run_ends().values();
1655        assert_eq!(run_ends.len(), 6);
1656        assert_eq!(&[2, 4, 5, 8, 9, 12], run_ends);
1657
1658        // Check the values
1659        let values = result_run_array
1660            .values()
1661            .as_any()
1662            .downcast_ref::<Int32Array>()
1663            .unwrap();
1664        assert_eq!(values.len(), 6);
1665        assert_eq!(&[10, 20, 30, 40, 50, 60], values.values());
1666    }
1667}