Skip to main content

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::{
41    ArrowNativeType, BooleanBufferBuilder, MutableBuffer, NullBuffer, OffsetBuffer, RunEndBuffer,
42    ScalarBuffer,
43};
44use arrow_data::transform::{Capacities, MutableArrayData};
45use arrow_schema::{ArrowError, DataType, FieldRef, Fields, SchemaRef};
46use std::{collections::HashSet, ops::Add, sync::Arc};
47
48fn binary_capacity<T: ByteArrayType>(arrays: &[&dyn Array]) -> Capacities {
49    let mut item_capacity = 0;
50    let mut bytes_capacity = 0;
51    for array in arrays {
52        let a = array.as_bytes::<T>();
53
54        // Guaranteed to always have at least one element
55        let offsets = a.value_offsets();
56        bytes_capacity += offsets[offsets.len() - 1].as_usize() - offsets[0].as_usize();
57        item_capacity += a.len()
58    }
59
60    Capacities::Binary(item_capacity, Some(bytes_capacity))
61}
62
63fn fixed_size_list_capacity(arrays: &[&dyn Array], data_type: &DataType) -> Capacities {
64    if let DataType::FixedSizeList(f, _) = data_type {
65        let item_capacity = arrays.iter().map(|a| a.len()).sum();
66        let child_data_type = f.data_type();
67        match child_data_type {
68            // These types should match the types that `get_capacity`
69            // has special handling for.
70            DataType::Utf8
71            | DataType::LargeUtf8
72            | DataType::Binary
73            | DataType::LargeBinary
74            | DataType::FixedSizeList(_, _) => {
75                let values: Vec<&dyn arrow_array::Array> = arrays
76                    .iter()
77                    .map(|a| a.as_fixed_size_list().values().as_ref())
78                    .collect();
79                Capacities::List(
80                    item_capacity,
81                    Some(Box::new(get_capacity(&values, child_data_type))),
82                )
83            }
84            _ => Capacities::Array(item_capacity),
85        }
86    } else {
87        unreachable!("illegal data type for fixed size list")
88    }
89}
90
91fn concat_byte_view<B: ByteViewType>(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError> {
92    let mut builder =
93        GenericByteViewBuilder::<B>::with_capacity(arrays.iter().map(|a| a.len()).sum());
94    for &array in arrays {
95        builder.append_array(array.as_byte_view());
96    }
97    Ok(Arc::new(builder.finish()))
98}
99
100fn concat_dictionaries<K: ArrowDictionaryKeyType>(
101    arrays: &[&dyn Array],
102) -> Result<ArrayRef, ArrowError> {
103    let mut output_len = 0;
104    let dictionaries: Vec<_> = arrays
105        .iter()
106        .map(|x| x.as_dictionary::<K>())
107        .inspect(|d| output_len += d.len())
108        .collect();
109
110    if !should_merge_dictionary_values::<K>(&dictionaries, output_len).0 {
111        return concat_fallback(arrays, Capacities::Array(output_len));
112    }
113
114    let merged = merge_dictionary_values(&dictionaries, None)?;
115
116    // Recompute keys
117    let mut key_values = Vec::with_capacity(output_len);
118
119    let mut has_nulls = false;
120    for (d, mapping) in dictionaries.iter().zip(merged.key_mappings) {
121        has_nulls |= d.null_count() != 0;
122        for key in d.keys().values() {
123            // Use get to safely handle nulls
124            key_values.push(mapping.get(key.as_usize()).copied().unwrap_or_default())
125        }
126    }
127
128    let nulls = has_nulls.then(|| {
129        let mut nulls = BooleanBufferBuilder::new(output_len);
130        for d in &dictionaries {
131            match d.nulls() {
132                Some(n) => nulls.append_buffer(n.inner()),
133                None => nulls.append_n(d.len(), true),
134            }
135        }
136        NullBuffer::new(nulls.finish())
137    });
138
139    let keys = PrimitiveArray::<K>::try_new(key_values.into(), nulls)?;
140    // Sanity check
141    assert_eq!(keys.len(), output_len);
142
143    let array = unsafe { DictionaryArray::new_unchecked(keys, merged.values) };
144    Ok(Arc::new(array))
145}
146
147/// Concatenates the child values and the validity bitmaps, without the per-array
148/// [`ArrayData`](arrow_data::ArrayData) that [`concat_fallback`] builds: with many short arrays
149/// that is most of the cost.
150fn concat_fixed_size_list(
151    arrays: &[&dyn Array],
152    field: &FieldRef,
153    size: i32,
154) -> Result<ArrayRef, ArrowError> {
155    let mut lists: Vec<&FixedSizeListArray> = Vec::with_capacity(arrays.len());
156    let mut values: Vec<&dyn Array> = Vec::with_capacity(arrays.len());
157    let (mut len, mut any_nulls) = (0, false);
158    for a in arrays {
159        let l = a.as_fixed_size_list();
160        len += l.len();
161        any_nulls |= l.null_count() != 0;
162        values.push(l.values().as_ref());
163        lists.push(l);
164    }
165
166    let nulls = any_nulls.then(|| {
167        let mut nulls = BooleanBufferBuilder::new(len);
168        for l in &lists {
169            match l.nulls() {
170                Some(n) => nulls.append_buffer(n.inner()),
171                None => nulls.append_n(l.len(), true),
172            }
173        }
174        NullBuffer::new(nulls.finish())
175    });
176
177    // Equal list types imply equal child types, so skip the check in `concat`.
178    let values = concat_same_type(&values, field.data_type())?;
179
180    Ok(Arc::new(FixedSizeListArray::try_new_with_length(
181        Arc::clone(field),
182        size,
183        values,
184        nulls,
185        len,
186    )?))
187}
188
189fn concat_lists<OffsetSize: OffsetSizeTrait>(
190    arrays: &[&dyn Array],
191    field: &FieldRef,
192) -> Result<ArrayRef, ArrowError> {
193    let mut output_len = 0;
194    let mut list_has_nulls = false;
195    let mut list_has_slices = false;
196
197    let lists = arrays
198        .iter()
199        .map(|x| x.as_list::<OffsetSize>())
200        .inspect(|l| {
201            output_len += l.len();
202            list_has_nulls |= l.null_count() != 0;
203            list_has_slices |= l.offsets()[0] > OffsetSize::zero()
204                || l.offsets().last().as_usize() < l.values().len();
205        })
206        .collect::<Vec<_>>();
207
208    let lists_nulls = list_has_nulls.then(|| {
209        let mut nulls = BooleanBufferBuilder::new(output_len);
210        for l in &lists {
211            match l.nulls() {
212                Some(n) => nulls.append_buffer(n.inner()),
213                None => nulls.append_n(l.len(), true),
214            }
215        }
216        NullBuffer::new(nulls.finish())
217    });
218
219    // If any of the lists have slices, we need to slice the values
220    // to ensure that the offsets are correct
221    let mut sliced_values;
222    let values: Vec<&dyn Array> = if list_has_slices {
223        sliced_values = Vec::with_capacity(lists.len());
224        for l in &lists {
225            // if the first offset is non-zero, we need to slice the values so when
226            // we concatenate them below only the relevant values are included
227            let offsets = l.offsets();
228            let start_offset = offsets[0].as_usize();
229            let end_offset = offsets.last().as_usize();
230            sliced_values.push(l.values().slice(start_offset, end_offset - start_offset));
231        }
232        sliced_values.iter().map(|a| a.as_ref()).collect()
233    } else {
234        lists.iter().map(|x| x.values().as_ref()).collect()
235    };
236
237    let concatenated_values = concat(values.as_slice())?;
238
239    // Merge value offsets from the lists
240    let value_offset_buffer =
241        OffsetBuffer::<OffsetSize>::from_lengths(lists.iter().flat_map(|x| x.offsets().lengths()));
242
243    let array = GenericListArray::<OffsetSize>::try_new(
244        Arc::clone(field),
245        value_offset_buffer,
246        concatenated_values,
247        lists_nulls,
248    )?;
249
250    Ok(Arc::new(array))
251}
252
253fn concat_maps(
254    arrays: &[&dyn Array],
255    field: &FieldRef,
256    ordered: bool,
257) -> Result<ArrayRef, ArrowError> {
258    let mut output_len = 0;
259    let mut map_has_nulls = false;
260    let mut map_has_slices = false;
261
262    let maps = arrays
263        .iter()
264        .map(|x| x.as_map())
265        .inspect(|m| {
266            output_len += m.len();
267            map_has_nulls |= m.null_count() != 0;
268            map_has_slices |=
269                m.offsets()[0] > 0 || m.offsets().last().as_usize() < m.entries().len();
270        })
271        .collect::<Vec<_>>();
272
273    let map_nulls = map_has_nulls.then(|| {
274        let mut nulls = BooleanBufferBuilder::new(output_len);
275        for m in &maps {
276            match m.nulls() {
277                Some(n) => nulls.append_buffer(n.inner()),
278                None => nulls.append_n(m.len(), true),
279            }
280        }
281        NullBuffer::new(nulls.finish())
282    });
283
284    // If any of the maps have slices, we need to slice the entries
285    // to ensure that the offsets are correct
286    let mut sliced_entries: Vec<ArrayRef>;
287    let entries: Vec<&dyn Array> = if map_has_slices {
288        sliced_entries = Vec::with_capacity(maps.len());
289        for m in &maps {
290            let offsets = m.offsets();
291            let start_offset = offsets[0].as_usize();
292            let end_offset = offsets.last().as_usize();
293            let entries_arr: &dyn Array = m.entries();
294            sliced_entries.push(entries_arr.slice(start_offset, end_offset - start_offset));
295        }
296        sliced_entries.iter().map(|a| a.as_ref()).collect()
297    } else {
298        maps.iter().map(|m| m.entries() as &dyn Array).collect()
299    };
300
301    let concatenated_entries = concat(entries.as_slice())?;
302
303    // Merge value offsets from the maps
304    let value_offset_buffer =
305        OffsetBuffer::<i32>::from_lengths(maps.iter().flat_map(|m| m.offsets().lengths()));
306
307    let array = MapArray::try_new(
308        Arc::clone(field),
309        value_offset_buffer,
310        // Safety: Map entries are always StructArrays, so this downcast is guaranteed to succeed
311        concatenated_entries.as_struct().clone(),
312        map_nulls,
313        ordered,
314    )?;
315
316    Ok(Arc::new(array))
317}
318
319fn concat_list_view<OffsetSize: OffsetSizeTrait>(
320    arrays: &[&dyn Array],
321    field: &FieldRef,
322) -> Result<ArrayRef, ArrowError> {
323    let mut output_len = 0;
324    let mut list_has_nulls = false;
325
326    let lists = arrays
327        .iter()
328        .map(|x| x.as_list_view::<OffsetSize>())
329        .inspect(|l| {
330            output_len += l.len();
331            list_has_nulls |= l.null_count() != 0;
332        })
333        .collect::<Vec<_>>();
334
335    let lists_nulls = list_has_nulls.then(|| {
336        let mut nulls = BooleanBufferBuilder::new(output_len);
337        for l in &lists {
338            match l.nulls() {
339                Some(n) => nulls.append_buffer(n.inner()),
340                None => nulls.append_n(l.len(), true),
341            }
342        }
343        NullBuffer::new(nulls.finish())
344    });
345
346    let values: Vec<&dyn Array> = lists.iter().map(|l| l.values().as_ref()).collect();
347
348    let concatenated_values = concat(values.as_slice())?;
349
350    let sizes: ScalarBuffer<OffsetSize> = lists.iter().flat_map(|x| x.sizes()).copied().collect();
351
352    let mut offsets = MutableBuffer::with_capacity(lists.iter().map(|l| l.offsets().len()).sum());
353    let mut global_offset = OffsetSize::zero();
354    for l in &lists {
355        for &offset in l.offsets() {
356            offsets.push(offset + global_offset);
357        }
358
359        // advance the offsets
360        global_offset += OffsetSize::from_usize(l.values().len()).unwrap();
361    }
362
363    let offsets = ScalarBuffer::from(offsets);
364
365    let array = GenericListViewArray::try_new(
366        field.clone(),
367        offsets,
368        sizes,
369        concatenated_values,
370        lists_nulls,
371    )?;
372
373    Ok(Arc::new(array))
374}
375
376fn concat_primitives<T: ArrowPrimitiveType>(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError> {
377    let mut builder = PrimitiveBuilder::<T>::with_capacity(arrays.iter().map(|a| a.len()).sum())
378        .with_data_type(arrays[0].data_type().clone());
379
380    for array in arrays {
381        builder.append_array(array.as_primitive());
382    }
383
384    Ok(Arc::new(builder.finish()))
385}
386
387fn concat_boolean(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError> {
388    let mut builder = BooleanBuilder::with_capacity(arrays.iter().map(|a| a.len()).sum());
389
390    for array in arrays {
391        builder.append_array(array.as_boolean());
392    }
393
394    Ok(Arc::new(builder.finish()))
395}
396
397fn concat_bytes<T: ByteArrayType>(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError> {
398    let Capacities::Binary(item_capacity, Some(bytes_capacity)) = binary_capacity::<T>(arrays)
399    else {
400        unreachable!()
401    };
402
403    let mut builder = GenericByteBuilder::<T>::with_capacity(item_capacity, bytes_capacity);
404
405    for array in arrays {
406        builder.append_array(array.as_bytes::<T>())?;
407    }
408
409    Ok(Arc::new(builder.finish()))
410}
411
412fn concat_structs(arrays: &[&dyn Array], fields: &Fields) -> Result<ArrayRef, ArrowError> {
413    let mut len = 0;
414    let mut has_nulls = false;
415    let structs = arrays
416        .iter()
417        .map(|a| {
418            len += a.len();
419            has_nulls |= a.null_count() > 0;
420            a.as_struct()
421        })
422        .collect::<Vec<_>>();
423
424    let nulls = has_nulls.then(|| {
425        let mut b = BooleanBufferBuilder::new(len);
426        for s in &structs {
427            match s.nulls() {
428                Some(n) => b.append_buffer(n.inner()),
429                None => b.append_n(s.len(), true),
430            }
431        }
432        NullBuffer::new(b.finish())
433    });
434
435    let column_concat_result = (0..fields.len())
436        .map(|i| {
437            let extracted_cols = structs
438                .iter()
439                .map(|s| s.column(i).as_ref())
440                .collect::<Vec<_>>();
441            concat(&extracted_cols)
442        })
443        .collect::<Result<Vec<_>, ArrowError>>()?;
444
445    Ok(Arc::new(StructArray::try_new_with_length(
446        fields.clone(),
447        column_concat_result,
448        nulls,
449        len,
450    )?))
451}
452
453/// Concatenate multiple RunArray instances into a single RunArray.
454///
455/// This function handles the special case of concatenating RunArrays by:
456/// 1. Collecting all run ends and values from input arrays
457/// 2. Adjusting run ends to account for the length of previous arrays
458/// 3. Creating a new RunArray with the combined data
459fn concat_run_arrays<R: RunEndIndexType>(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError>
460where
461    R::Native: Add<Output = R::Native>,
462{
463    let run_arrays: Vec<_> = arrays
464        .iter()
465        .map(|x| x.as_run::<R>())
466        .filter(|x| !x.run_ends().is_empty())
467        .collect();
468
469    if run_arrays.is_empty() {
470        // If all input arrays are empty then handle here otherwise we
471        // lose the type below
472        return Ok(new_empty_array(arrays[0].data_type()));
473    }
474
475    // Reject lengths that do not fit in `R` before any `R::Native` arithmetic can wrap.
476    let total_len: usize = run_arrays.iter().map(|r| r.len()).sum();
477    if R::Native::from_usize(total_len).is_none() {
478        return Err(ArrowError::RunEndIndexOverflowError);
479    }
480
481    // The run ends need to be adjusted by the sum of the lengths of the previous arrays.
482    let needed_run_end_adjustments = std::iter::once(R::default_value())
483        .chain(
484            run_arrays
485                .iter()
486                .scan(R::default_value(), |acc, run_array| {
487                    *acc = *acc + R::Native::from_usize(run_array.len()).unwrap();
488                    Some(*acc)
489                }),
490        )
491        .collect::<Vec<_>>();
492
493    let run_ends_array =
494        PrimitiveArray::<R>::from_iter_values(run_arrays.iter().enumerate().flat_map(
495            move |(i, run_array)| {
496                let adjustment = needed_run_end_adjustments[i];
497                run_array
498                    .run_ends()
499                    .sliced_values()
500                    .map(move |run_end| run_end + adjustment)
501            },
502        ));
503
504    let values_slices: Vec<ArrayRef> = run_arrays
505        .iter()
506        .map(|run_array| run_array.values_slice())
507        .collect();
508
509    let all_values = concat(&values_slices.iter().map(|x| x.as_ref()).collect::<Vec<_>>())?;
510
511    let data_type = run_arrays[0].data_type().clone();
512    let (_, run_ends_values, _) = run_ends_array.into_parts();
513
514    // Safety: inputs are valid RunArrays; adjusted run ends are strictly increasing
515    // and end at `total_len`. Physical length matches `all_values`.
516    let run_ends = unsafe { RunEndBuffer::new_unchecked(run_ends_values, 0, total_len) };
517    Ok(Arc::new(unsafe {
518        RunArray::<R>::new_unchecked(data_type, run_ends, all_values)
519    }))
520}
521
522macro_rules! dict_helper {
523    ($t:ty, $arrays:expr) => {
524        return concat_dictionaries::<$t>($arrays)
525    };
526}
527
528macro_rules! primitive_concat {
529    ($t:ty, $arrays:expr) => {
530        return concat_primitives::<$t>($arrays)
531    };
532}
533
534fn get_capacity(arrays: &[&dyn Array], data_type: &DataType) -> Capacities {
535    match data_type {
536        DataType::Utf8 => binary_capacity::<Utf8Type>(arrays),
537        DataType::LargeUtf8 => binary_capacity::<LargeUtf8Type>(arrays),
538        DataType::Binary => binary_capacity::<BinaryType>(arrays),
539        DataType::LargeBinary => binary_capacity::<LargeBinaryType>(arrays),
540        DataType::FixedSizeList(_, _) => fixed_size_list_capacity(arrays, data_type),
541        _ => Capacities::Array(arrays.iter().map(|a| a.len()).sum()),
542    }
543}
544
545/// Concatenate multiple [Array] of the same type into a single [ArrayRef].
546pub fn concat(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError> {
547    if arrays.is_empty() {
548        return Err(ArrowError::ComputeError(
549            "concat requires input of at least one array".to_string(),
550        ));
551    } else if arrays.len() == 1 {
552        let array = arrays[0];
553        return Ok(array.slice(0, array.len()));
554    }
555
556    let d = arrays[0].data_type();
557    if arrays.iter().skip(1).any(|array| array.data_type() != d) {
558        // Create error message with up to 10 unique data types in the order they appear
559        let error_message = {
560            // 10 max unique data types to print and another 1 to know if there are more
561            let mut unique_data_types = HashSet::with_capacity(11);
562
563            let mut error_message =
564                format!("It is not possible to concatenate arrays of different data types ({d}");
565            unique_data_types.insert(d);
566
567            for array in arrays {
568                let is_unique = unique_data_types.insert(array.data_type());
569
570                if unique_data_types.len() == 11 {
571                    error_message.push_str(", ...");
572                    break;
573                }
574
575                if is_unique {
576                    error_message.push_str(", ");
577                    error_message.push_str(&array.data_type().to_string());
578                }
579            }
580
581            error_message.push_str(").");
582
583            error_message
584        };
585
586        return Err(ArrowError::InvalidArgumentError(error_message));
587    }
588
589    concat_same_type(arrays, d)
590}
591
592fn concat_same_type(arrays: &[&dyn Array], d: &DataType) -> Result<ArrayRef, ArrowError> {
593    downcast_primitive! {
594        d => (primitive_concat, arrays),
595        DataType::Boolean => concat_boolean(arrays),
596        DataType::Dictionary(k, _) => {
597            downcast_integer! {
598                k.as_ref() => (dict_helper, arrays),
599                _ => unreachable!("illegal dictionary key type {k}")
600            }
601        }
602        DataType::List(field) => concat_lists::<i32>(arrays, field),
603        DataType::LargeList(field) => concat_lists::<i64>(arrays, field),
604        DataType::FixedSizeList(field, size) => concat_fixed_size_list(arrays, field, *size),
605        DataType::ListView(field) => concat_list_view::<i32>(arrays, field),
606        DataType::LargeListView(field) => concat_list_view::<i64>(arrays, field),
607        DataType::Map(field, ordered) => concat_maps(arrays, field, *ordered),
608        DataType::Struct(fields) => concat_structs(arrays, fields),
609        DataType::Utf8 => concat_bytes::<Utf8Type>(arrays),
610        DataType::LargeUtf8 => concat_bytes::<LargeUtf8Type>(arrays),
611        DataType::Binary => concat_bytes::<BinaryType>(arrays),
612        DataType::LargeBinary => concat_bytes::<LargeBinaryType>(arrays),
613        DataType::RunEndEncoded(r, _) => {
614            // Handle RunEndEncoded arrays with special concat function
615            // We need to downcast based on the run end type
616            match r.data_type() {
617                DataType::Int16 => concat_run_arrays::<Int16Type>(arrays),
618                DataType::Int32 => concat_run_arrays::<Int32Type>(arrays),
619                DataType::Int64 => concat_run_arrays::<Int64Type>(arrays),
620                _ => unreachable!("Unsupported run end index type: {r:?}"),
621            }
622        }
623        DataType::Utf8View => concat_byte_view::<StringViewType>(arrays),
624        DataType::BinaryView => concat_byte_view::<BinaryViewType>(arrays),
625        _ => {
626            let capacity = get_capacity(arrays, d);
627            concat_fallback(arrays, capacity)
628        }
629    }
630}
631
632/// Concatenates arrays using MutableArrayData
633///
634/// This will naively concatenate dictionaries
635fn concat_fallback(arrays: &[&dyn Array], capacity: Capacities) -> Result<ArrayRef, ArrowError> {
636    let array_data: Vec<_> = arrays.iter().map(|a| a.to_data()).collect::<Vec<_>>();
637    let array_data = array_data.iter().collect();
638    let mut mutable = MutableArrayData::try_with_capacities(array_data, false, capacity)?;
639
640    for (i, a) in arrays.iter().enumerate() {
641        mutable.try_extend(i, 0, a.len())?
642    }
643
644    Ok(make_array(mutable.freeze()))
645}
646
647/// Concatenates `batches` together into a single [`RecordBatch`].
648///
649/// The output batch has the specified `schemas`; The schema of the
650/// input are ignored.
651///
652/// # Notes
653///
654/// - Callers should budget for peak memory use to approach 2x the input
655///   size, as the input batches and output arrays co-exist during construction.
656/// - Arrays with `i32` offsets, such as `StringArray` and `BinaryArray`, only
657///   support up to ~2GiB of payloads. Concatenating large arrays of these types
658///   can cause offset overflows.
659///
660/// # Errors
661///
662/// Returns an error if the types of underlying arrays are different.
663pub fn concat_batches<'a>(
664    schema: &SchemaRef,
665    input_batches: impl IntoIterator<Item = &'a RecordBatch>,
666) -> Result<RecordBatch, ArrowError> {
667    // When schema is empty, sum the number of the rows of all batches
668    if schema.fields().is_empty() {
669        let num_rows: usize = input_batches.into_iter().map(RecordBatch::num_rows).sum();
670        let mut options = RecordBatchOptions::default();
671        options.row_count = Some(num_rows);
672        return RecordBatch::try_new_with_options(schema.clone(), vec![], &options);
673    }
674
675    let batches: Vec<&RecordBatch> = input_batches.into_iter().collect();
676    if batches.is_empty() {
677        return Ok(RecordBatch::new_empty(schema.clone()));
678    }
679    let field_num = schema.fields().len();
680    let mut arrays = Vec::with_capacity(field_num);
681    for i in 0..field_num {
682        let array = concat(
683            &batches
684                .iter()
685                .map(|batch| batch.column(i).as_ref())
686                .collect::<Vec<_>>(),
687        )?;
688        arrays.push(array);
689    }
690    RecordBatch::try_new(schema.clone(), arrays)
691}
692
693#[cfg(test)]
694mod tests {
695    use super::*;
696    use arrow_array::builder::{
697        GenericListBuilder, Int32Builder as Int32ArrayBuilder, Int64Builder, ListViewBuilder,
698        MapBuilder, StringBuilder, StringDictionaryBuilder,
699    };
700    use arrow_schema::{Field, Schema};
701    use std::fmt::Debug;
702
703    #[test]
704    fn test_dict_overflow_9366() {
705        use arrow_schema::DataType;
706
707        let schema = Arc::new(Schema::new(vec![Field::new(
708            "a",
709            DataType::Dictionary(
710                Box::new(DataType::UInt8),
711                Box::new(DataType::FixedSizeBinary(8)),
712            ),
713            false,
714        )]));
715        let make = |vals: std::ops::Range<u64>| {
716            let dict = FixedSizeBinaryArray::try_from_iter(vals.map(|i| i.to_le_bytes())).unwrap();
717            let keys = UInt8Array::from_iter_values(0..128);
718            let arr = DictionaryArray::try_new(keys, Arc::new(dict)).unwrap();
719            RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap()
720        };
721        // 256 distinct values fit in u8 keys (0..=255): concat must succeed.
722        let out = concat_batches(&schema, &[make(0..128), make(128..256)]).unwrap();
723        assert_eq!(out.num_rows(), 256);
724        let dict = out.column(0).as_dictionary::<UInt8Type>();
725        assert_eq!(dict.values().len(), 256);
726    }
727
728    #[test]
729    fn test_dict_overflow_i8_9366() {
730        use arrow_schema::DataType;
731
732        // Same boundary for a signed key type: i8 holds 128 keys (0..=127).
733        let schema = Arc::new(Schema::new(vec![Field::new(
734            "a",
735            DataType::Dictionary(
736                Box::new(DataType::Int8),
737                Box::new(DataType::FixedSizeBinary(8)),
738            ),
739            false,
740        )]));
741        let make = |vals: std::ops::Range<u64>| {
742            let dict = FixedSizeBinaryArray::try_from_iter(vals.map(|i| i.to_le_bytes())).unwrap();
743            let keys = Int8Array::from_iter_values(0..64);
744            let arr = DictionaryArray::try_new(keys, Arc::new(dict)).unwrap();
745            RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap()
746        };
747        let out = concat_batches(&schema, &[make(0..64), make(64..128)]).unwrap();
748        assert_eq!(out.num_rows(), 128);
749        let dict = out.column(0).as_dictionary::<Int8Type>();
750        assert_eq!(dict.values().len(), 128);
751    }
752
753    #[test]
754    fn test_concat_empty_vec() {
755        let re = concat(&[]);
756        assert!(re.is_err());
757    }
758
759    #[test]
760    fn test_concat_batches_no_columns() {
761        // Test concat using empty schema / batches without columns
762        let schema = Arc::new(Schema::empty());
763
764        let mut options = RecordBatchOptions::default();
765        options.row_count = Some(100);
766        let batch = RecordBatch::try_new_with_options(schema.clone(), vec![], &options).unwrap();
767        // put in 2 batches of 100 rows each
768        let re = concat_batches(&schema, &[batch.clone(), batch]).unwrap();
769
770        assert_eq!(re.num_rows(), 200);
771    }
772
773    #[test]
774    fn test_concat_one_element_vec() {
775        let arr = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
776            Some(-1),
777            Some(2),
778            None,
779        ])) as ArrayRef;
780        let result = concat(&[arr.as_ref()]).unwrap();
781        assert_eq!(
782            &arr, &result,
783            "concatenating single element array gives back the same result"
784        );
785    }
786
787    #[test]
788    fn test_concat_incompatible_datatypes() {
789        let re = concat(&[
790            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(2), None]),
791            // 2 string to make sure we only mention unique types
792            &StringArray::from(vec![Some("hello"), Some("bar"), Some("world")]),
793            &StringArray::from(vec![Some("hey"), Some(""), Some("you")]),
794            // Another type to make sure we are showing all the incompatible types
795            &PrimitiveArray::<Int32Type>::from(vec![Some(-1), Some(2), None]),
796        ]);
797
798        assert_eq!(
799            re.unwrap_err().to_string(),
800            "Invalid argument error: It is not possible to concatenate arrays of different data types (Int64, Utf8, Int32)."
801        );
802    }
803
804    #[test]
805    fn test_concat_10_incompatible_datatypes_should_include_all_of_them() {
806        let re = concat(&[
807            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(2), None]),
808            // 2 string to make sure we only mention unique types
809            &StringArray::from(vec![Some("hello"), Some("bar"), Some("world")]),
810            &StringArray::from(vec![Some("hey"), Some(""), Some("you")]),
811            // Another type to make sure we are showing all the incompatible types
812            &PrimitiveArray::<Int32Type>::from(vec![Some(-1), Some(2), None]),
813            &PrimitiveArray::<Int8Type>::from(vec![Some(-1), Some(2), None]),
814            &PrimitiveArray::<Int16Type>::from(vec![Some(-1), Some(2), None]),
815            &PrimitiveArray::<UInt8Type>::from(vec![Some(1), Some(2), None]),
816            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
817            &PrimitiveArray::<UInt32Type>::from(vec![Some(1), Some(2), None]),
818            // Non unique
819            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
820            &PrimitiveArray::<UInt64Type>::from(vec![Some(1), Some(2), None]),
821            &PrimitiveArray::<Float32Type>::from(vec![Some(1.0), Some(2.0), None]),
822        ]);
823
824        assert_eq!(
825            re.unwrap_err().to_string(),
826            "Invalid argument error: It is not possible to concatenate arrays of different data types (Int64, Utf8, Int32, Int8, Int16, UInt8, UInt16, UInt32, UInt64, Float32)."
827        );
828    }
829
830    #[test]
831    fn test_concat_11_incompatible_datatypes_should_only_include_10() {
832        let re = concat(&[
833            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(2), None]),
834            // 2 string to make sure we only mention unique types
835            &StringArray::from(vec![Some("hello"), Some("bar"), Some("world")]),
836            &StringArray::from(vec![Some("hey"), Some(""), Some("you")]),
837            // Another type to make sure we are showing all the incompatible types
838            &PrimitiveArray::<Int32Type>::from(vec![Some(-1), Some(2), None]),
839            &PrimitiveArray::<Int8Type>::from(vec![Some(-1), Some(2), None]),
840            &PrimitiveArray::<Int16Type>::from(vec![Some(-1), Some(2), None]),
841            &PrimitiveArray::<UInt8Type>::from(vec![Some(1), Some(2), None]),
842            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
843            &PrimitiveArray::<UInt32Type>::from(vec![Some(1), Some(2), None]),
844            // Non unique
845            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
846            &PrimitiveArray::<UInt64Type>::from(vec![Some(1), Some(2), None]),
847            &PrimitiveArray::<Float32Type>::from(vec![Some(1.0), Some(2.0), None]),
848            &PrimitiveArray::<Float64Type>::from(vec![Some(1.0), Some(2.0), None]),
849        ]);
850
851        assert_eq!(
852            re.unwrap_err().to_string(),
853            "Invalid argument error: It is not possible to concatenate arrays of different data types (Int64, Utf8, Int32, Int8, Int16, UInt8, UInt16, UInt32, UInt64, Float32, ...)."
854        );
855    }
856
857    #[test]
858    #[cfg_attr(miri, ignore)] // Unsupported inline assembly
859    fn test_concat_13_incompatible_datatypes_should_not_include_all_of_them() {
860        let re = concat(&[
861            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(2), None]),
862            // 2 string to make sure we only mention unique types
863            &StringArray::from(vec![Some("hello"), Some("bar"), Some("world")]),
864            &StringArray::from(vec![Some("hey"), Some(""), Some("you")]),
865            // Another type to make sure we are showing all the incompatible types
866            &PrimitiveArray::<Int32Type>::from(vec![Some(-1), Some(2), None]),
867            &PrimitiveArray::<Int8Type>::from(vec![Some(-1), Some(2), None]),
868            &PrimitiveArray::<Int16Type>::from(vec![Some(-1), Some(2), None]),
869            &PrimitiveArray::<UInt8Type>::from(vec![Some(1), Some(2), None]),
870            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
871            &PrimitiveArray::<UInt32Type>::from(vec![Some(1), Some(2), None]),
872            // Non unique
873            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
874            &PrimitiveArray::<UInt64Type>::from(vec![Some(1), Some(2), None]),
875            &PrimitiveArray::<Float32Type>::from(vec![Some(1.0), Some(2.0), None]),
876            &PrimitiveArray::<Float64Type>::from(vec![Some(1.0), Some(2.0), None]),
877            &PrimitiveArray::<Float16Type>::new_null(3),
878            &BooleanArray::from(vec![Some(true), Some(false), None]),
879        ]);
880
881        assert_eq!(
882            re.unwrap_err().to_string(),
883            "Invalid argument error: It is not possible to concatenate arrays of different data types (Int64, Utf8, Int32, Int8, Int16, UInt8, UInt16, UInt32, UInt64, Float32, ...)."
884        );
885    }
886
887    #[test]
888    fn test_concat_string_arrays() {
889        let arr = concat(&[
890            &StringArray::from(vec!["hello", "world"]),
891            &StringArray::from(vec!["2", "3", "4"]),
892            &StringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]),
893        ])
894        .unwrap();
895
896        let expected_output = Arc::new(StringArray::from(vec![
897            Some("hello"),
898            Some("world"),
899            Some("2"),
900            Some("3"),
901            Some("4"),
902            Some("foo"),
903            Some("bar"),
904            None,
905            Some("baz"),
906        ])) as ArrayRef;
907
908        assert_eq!(&arr, &expected_output);
909    }
910
911    #[test]
912    fn test_concat_string_view_arrays() {
913        let arr = concat(&[
914            &StringViewArray::from(vec!["helloxxxxxxxxxxa", "world____________"]),
915            &StringViewArray::from(vec!["helloxxxxxxxxxxy", "3", "4"]),
916            &StringViewArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]),
917        ])
918        .unwrap();
919
920        let expected_output = Arc::new(StringViewArray::from(vec![
921            Some("helloxxxxxxxxxxa"),
922            Some("world____________"),
923            Some("helloxxxxxxxxxxy"),
924            Some("3"),
925            Some("4"),
926            Some("foo"),
927            Some("bar"),
928            None,
929            Some("baz"),
930        ])) as ArrayRef;
931
932        assert_eq!(&arr, &expected_output);
933    }
934
935    #[test]
936    fn test_concat_primitive_arrays() {
937        let arr = concat(&[
938            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(-1), Some(2), None, None]),
939            &PrimitiveArray::<Int64Type>::from(vec![Some(101), Some(102), Some(103), None]),
940            &PrimitiveArray::<Int64Type>::from(vec![Some(256), Some(512), Some(1024)]),
941        ])
942        .unwrap();
943
944        let expected_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
945            Some(-1),
946            Some(-1),
947            Some(2),
948            None,
949            None,
950            Some(101),
951            Some(102),
952            Some(103),
953            None,
954            Some(256),
955            Some(512),
956            Some(1024),
957        ])) as ArrayRef;
958
959        assert_eq!(&arr, &expected_output);
960    }
961
962    #[test]
963    fn test_concat_primitive_array_slices() {
964        let input_1 =
965            PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(-1), Some(2), None, None])
966                .slice(1, 3);
967
968        let input_2 =
969            PrimitiveArray::<Int64Type>::from(vec![Some(101), Some(102), Some(103), None])
970                .slice(1, 3);
971        let arr = concat(&[&input_1, &input_2]).unwrap();
972
973        let expected_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
974            Some(-1),
975            Some(2),
976            None,
977            Some(102),
978            Some(103),
979            None,
980        ])) as ArrayRef;
981
982        assert_eq!(&arr, &expected_output);
983    }
984
985    #[test]
986    fn test_concat_boolean_primitive_arrays() {
987        let arr = concat(&[
988            &BooleanArray::from(vec![
989                Some(true),
990                Some(true),
991                Some(false),
992                None,
993                None,
994                Some(false),
995            ]),
996            &BooleanArray::from(vec![None, Some(false), Some(true), Some(false)]),
997        ])
998        .unwrap();
999
1000        let expected_output = Arc::new(BooleanArray::from(vec![
1001            Some(true),
1002            Some(true),
1003            Some(false),
1004            None,
1005            None,
1006            Some(false),
1007            None,
1008            Some(false),
1009            Some(true),
1010            Some(false),
1011        ])) as ArrayRef;
1012
1013        assert_eq!(&arr, &expected_output);
1014    }
1015
1016    #[test]
1017    fn test_concat_primitive_list_arrays() {
1018        let list1 = [
1019            Some(vec![Some(-1), Some(-1), Some(2), None, None]),
1020            Some(vec![]),
1021            None,
1022            Some(vec![Some(10)]),
1023        ];
1024        let list1_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone());
1025
1026        let list2 = [
1027            None,
1028            Some(vec![Some(100), None, Some(101)]),
1029            Some(vec![Some(102)]),
1030        ];
1031        let list2_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone());
1032
1033        let list3 = [Some(vec![Some(1000), Some(1001)])];
1034        let list3_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list3.clone());
1035
1036        let array_result = concat(&[&list1_array, &list2_array, &list3_array]).unwrap();
1037
1038        let expected = list1.into_iter().chain(list2).chain(list3);
1039        let array_expected = ListArray::from_iter_primitive::<Int64Type, _, _>(expected);
1040
1041        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1042    }
1043
1044    #[test]
1045    fn test_concat_primitive_list_arrays_slices() {
1046        let list1 = [
1047            Some(vec![Some(-1), Some(-1), Some(2), None, None]),
1048            Some(vec![]), // In slice
1049            None,         // In slice
1050            Some(vec![Some(10)]),
1051        ];
1052        let list1_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone());
1053        let list1_array = list1_array.slice(1, 2);
1054        let list1_values = list1.into_iter().skip(1).take(2);
1055
1056        let list2 = [
1057            None,
1058            Some(vec![Some(100), None, Some(101)]),
1059            Some(vec![Some(102)]),
1060        ];
1061        let list2_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone());
1062
1063        // verify that this test covers the case when the first offset is non zero
1064        assert!(list1_array.offsets()[0].as_usize() > 0);
1065        let array_result = concat(&[&list1_array, &list2_array]).unwrap();
1066
1067        let expected = list1_values.chain(list2);
1068        let array_expected = ListArray::from_iter_primitive::<Int64Type, _, _>(expected);
1069
1070        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1071    }
1072
1073    #[test]
1074    fn test_concat_primitive_list_arrays_sliced_lengths() {
1075        let list1 = [
1076            Some(vec![Some(-1), Some(-1), Some(2), None, None]), // In slice
1077            Some(vec![]),                                        // In slice
1078            None,                                                // In slice
1079            Some(vec![Some(10)]),
1080        ];
1081        let list1_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone());
1082        let list1_array = list1_array.slice(0, 3); // no offset, but not all values
1083        let list1_values = list1.into_iter().take(3);
1084
1085        let list2 = [
1086            None,
1087            Some(vec![Some(100), None, Some(101)]),
1088            Some(vec![Some(102)]),
1089        ];
1090        let list2_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone());
1091
1092        // verify that this test covers the case when the first offset is zero, but the
1093        // last offset doesn't cover the entire array
1094        assert_eq!(list1_array.offsets()[0].as_usize(), 0);
1095        assert!(list1_array.offsets().last().as_usize() < list1_array.values().len());
1096        let array_result = concat(&[&list1_array, &list2_array]).unwrap();
1097
1098        let expected = list1_values.chain(list2);
1099        let array_expected = ListArray::from_iter_primitive::<Int64Type, _, _>(expected);
1100
1101        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1102    }
1103
1104    #[test]
1105    fn test_concat_primitive_fixed_size_list_arrays() {
1106        let list1 = [
1107            Some(vec![Some(-1), None]),
1108            None,
1109            Some(vec![Some(10), Some(20)]),
1110        ];
1111        let list1_array =
1112            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone(), 2);
1113
1114        let list2 = [
1115            None,
1116            Some(vec![Some(100), None]),
1117            Some(vec![Some(102), Some(103)]),
1118        ];
1119        let list2_array =
1120            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone(), 2);
1121
1122        let list3 = [Some(vec![Some(1000), Some(1001)])];
1123        let list3_array =
1124            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(list3.clone(), 2);
1125
1126        let array_result = concat(&[&list1_array, &list2_array, &list3_array]).unwrap();
1127
1128        let expected = list1.into_iter().chain(list2).chain(list3);
1129        let array_expected =
1130            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(expected, 2);
1131
1132        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1133    }
1134
1135    /// The fixed size list path against `concat_fallback`: slices, arrays with and without
1136    /// nulls mixed, nested lists, zero-sized lists and many one-row arrays.
1137    #[test]
1138    fn test_concat_fixed_size_list_matches_fallback() {
1139        let check = |arrays: &[&dyn Array]| {
1140            let expected =
1141                concat_fallback(arrays, get_capacity(arrays, arrays[0].data_type())).unwrap();
1142            let got = concat(arrays).unwrap();
1143            assert_eq!(got.as_ref(), expected.as_ref());
1144            got.to_data().validate_full().unwrap();
1145        };
1146
1147        let rows = [
1148            Some(vec![Some(1), None, Some(3)]),
1149            None,
1150            Some(vec![Some(4), Some(5), Some(6)]),
1151            Some(vec![None, None, None]),
1152            None,
1153            Some(vec![Some(7), Some(8), Some(9)]),
1154        ];
1155        let with_nulls = FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>(rows, 3);
1156        let without_nulls = FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>(
1157            (0..5).map(|i| Some(vec![Some(i), Some(i + 1), Some(i + 2)])),
1158            3,
1159        );
1160        check(&[&with_nulls, &without_nulls]);
1161        check(&[&without_nulls, &without_nulls]);
1162        check(&[
1163            &with_nulls.slice(1, 4),
1164            &without_nulls.slice(2, 2),
1165            &with_nulls,
1166        ]);
1167
1168        let one_row: Vec<_> = (0..64).map(|i| with_nulls.slice(i % 6, 1)).collect();
1169        let one_row: Vec<&dyn Array> = one_row.iter().map(|a| a as &dyn Array).collect();
1170        check(&one_row);
1171
1172        let nested = FixedSizeListArray::try_new(
1173            Arc::new(Field::new_list_field(with_nulls.data_type().clone(), true)),
1174            2,
1175            Arc::new(with_nulls.slice(0, 6)),
1176            Some(NullBuffer::from(vec![true, false, true])),
1177        )
1178        .unwrap();
1179        check(&[&nested, &nested.slice(1, 2)]);
1180
1181        let zero_sized = FixedSizeListArray::try_new_with_length(
1182            Arc::new(Field::new_list_field(DataType::Int32, true)),
1183            0,
1184            Arc::new(Int32Array::from(Vec::<i32>::new())),
1185            Some(NullBuffer::from(vec![true, false, true, true])),
1186            4,
1187        )
1188        .unwrap();
1189        let zero_sized_valid = FixedSizeListArray::try_new_with_length(
1190            Arc::new(Field::new_list_field(DataType::Int32, true)),
1191            0,
1192            Arc::new(Int32Array::from(Vec::<i32>::new())),
1193            None,
1194            3,
1195        )
1196        .unwrap();
1197        check(&[&zero_sized, &zero_sized_valid, &zero_sized.slice(1, 2)]);
1198    }
1199
1200    #[test]
1201    fn test_concat_list_view_arrays() {
1202        let list1 = [
1203            Some(vec![Some(-1), None]),
1204            None,
1205            Some(vec![Some(10), Some(20)]),
1206        ];
1207        let mut list1_array = ListViewBuilder::new(Int64Builder::new());
1208        for v in &list1 {
1209            list1_array.append_option(v.clone());
1210        }
1211        let list1_array = list1_array.finish();
1212
1213        let list2 = [
1214            None,
1215            Some(vec![Some(100), None]),
1216            Some(vec![Some(102), Some(103)]),
1217        ];
1218        let mut list2_array = ListViewBuilder::new(Int64Builder::new());
1219        for v in &list2 {
1220            list2_array.append_option(v.clone());
1221        }
1222        let list2_array = list2_array.finish();
1223
1224        let list3 = [Some(vec![Some(1000), Some(1001)])];
1225        let mut list3_array = ListViewBuilder::new(Int64Builder::new());
1226        for v in &list3 {
1227            list3_array.append_option(v.clone());
1228        }
1229        let list3_array = list3_array.finish();
1230
1231        let array_result = concat(&[&list1_array, &list2_array, &list3_array]).unwrap();
1232
1233        let expected: Vec<_> = list1.into_iter().chain(list2).chain(list3).collect();
1234        let mut array_expected = ListViewBuilder::new(Int64Builder::new());
1235        for v in &expected {
1236            array_expected.append_option(v.clone());
1237        }
1238        let array_expected = array_expected.finish();
1239
1240        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1241    }
1242
1243    #[test]
1244    fn test_concat_sliced_list_view_arrays() {
1245        let list1 = [
1246            Some(vec![Some(-1), None]),
1247            None,
1248            Some(vec![Some(10), Some(20)]),
1249        ];
1250        let mut list1_array = ListViewBuilder::new(Int64Builder::new());
1251        for v in &list1 {
1252            list1_array.append_option(v.clone());
1253        }
1254        let list1_array = list1_array.finish();
1255
1256        let list2 = [
1257            None,
1258            Some(vec![Some(100), None]),
1259            Some(vec![Some(102), Some(103)]),
1260        ];
1261        let mut list2_array = ListViewBuilder::new(Int64Builder::new());
1262        for v in &list2 {
1263            list2_array.append_option(v.clone());
1264        }
1265        let list2_array = list2_array.finish();
1266
1267        let list3 = [Some(vec![Some(1000), Some(1001)])];
1268        let mut list3_array = ListViewBuilder::new(Int64Builder::new());
1269        for v in &list3 {
1270            list3_array.append_option(v.clone());
1271        }
1272        let list3_array = list3_array.finish();
1273
1274        // Concat sliced arrays.
1275        // ListView slicing will slice the offset/sizes but preserve the original values child.
1276        let array_result = concat(&[
1277            &list1_array.slice(1, 2),
1278            &list2_array.slice(1, 2),
1279            &list3_array.slice(0, 1),
1280        ])
1281        .unwrap();
1282
1283        let expected: Vec<_> = vec![
1284            None,
1285            Some(vec![Some(10), Some(20)]),
1286            Some(vec![Some(100), None]),
1287            Some(vec![Some(102), Some(103)]),
1288            Some(vec![Some(1000), Some(1001)]),
1289        ];
1290        let mut array_expected = ListViewBuilder::new(Int64Builder::new());
1291        for v in &expected {
1292            array_expected.append_option(v.clone());
1293        }
1294        let array_expected = array_expected.finish();
1295
1296        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1297    }
1298
1299    #[test]
1300    fn test_concat_struct_arrays() {
1301        let field = Arc::new(Field::new("field", DataType::Int64, true));
1302        let input_primitive_1: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1303            Some(-1),
1304            Some(-1),
1305            Some(2),
1306            None,
1307            None,
1308        ]));
1309        let input_struct_1 = StructArray::from(vec![(field.clone(), input_primitive_1)]);
1310
1311        let input_primitive_2: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1312            Some(101),
1313            Some(102),
1314            Some(103),
1315            None,
1316        ]));
1317        let input_struct_2 = StructArray::from(vec![(field.clone(), input_primitive_2)]);
1318
1319        let input_primitive_3: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1320            Some(256),
1321            Some(512),
1322            Some(1024),
1323        ]));
1324        let input_struct_3 = StructArray::from(vec![(field, input_primitive_3)]);
1325
1326        let arr = concat(&[&input_struct_1, &input_struct_2, &input_struct_3]).unwrap();
1327
1328        let expected_primitive_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1329            Some(-1),
1330            Some(-1),
1331            Some(2),
1332            None,
1333            None,
1334            Some(101),
1335            Some(102),
1336            Some(103),
1337            None,
1338            Some(256),
1339            Some(512),
1340            Some(1024),
1341        ])) as ArrayRef;
1342
1343        let actual_primitive = arr
1344            .as_any()
1345            .downcast_ref::<StructArray>()
1346            .unwrap()
1347            .column(0);
1348        assert_eq!(actual_primitive, &expected_primitive_output);
1349    }
1350
1351    #[test]
1352    fn test_concat_struct_array_slices() {
1353        let field = Arc::new(Field::new("field", DataType::Int64, true));
1354        let input_primitive_1: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1355            Some(-1),
1356            Some(-1),
1357            Some(2),
1358            None,
1359            None,
1360        ]));
1361        let input_struct_1 = StructArray::from(vec![(field.clone(), input_primitive_1)]);
1362
1363        let input_primitive_2: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1364            Some(101),
1365            Some(102),
1366            Some(103),
1367            None,
1368        ]));
1369        let input_struct_2 = StructArray::from(vec![(field, input_primitive_2)]);
1370
1371        let arr = concat(&[&input_struct_1.slice(1, 3), &input_struct_2.slice(1, 2)]).unwrap();
1372
1373        let expected_primitive_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1374            Some(-1),
1375            Some(2),
1376            None,
1377            Some(102),
1378            Some(103),
1379        ])) as ArrayRef;
1380
1381        let actual_primitive = arr
1382            .as_any()
1383            .downcast_ref::<StructArray>()
1384            .unwrap()
1385            .column(0);
1386        assert_eq!(actual_primitive, &expected_primitive_output);
1387    }
1388
1389    #[test]
1390    fn test_concat_struct_arrays_no_nulls() {
1391        let input_1a = vec![1, 2, 3];
1392        let input_1b = vec!["one", "two", "three"];
1393        let input_2a = vec![4, 5, 6, 7];
1394        let input_2b = vec!["four", "five", "six", "seven"];
1395
1396        let struct_from_primitives = |ints: Vec<i64>, strings: Vec<&str>| {
1397            StructArray::try_from(vec![
1398                ("ints", Arc::new(Int64Array::from(ints)) as _),
1399                ("strings", Arc::new(StringArray::from(strings)) as _),
1400            ])
1401        };
1402
1403        let expected_output = struct_from_primitives(
1404            [input_1a.clone(), input_2a.clone()].concat(),
1405            [input_1b.clone(), input_2b.clone()].concat(),
1406        )
1407        .unwrap();
1408
1409        let input_1 = struct_from_primitives(input_1a, input_1b).unwrap();
1410        let input_2 = struct_from_primitives(input_2a, input_2b).unwrap();
1411
1412        let arr = concat(&[&input_1, &input_2]).unwrap();
1413        let struct_result = arr.as_struct();
1414
1415        assert_eq!(struct_result, &expected_output);
1416        assert_eq!(arr.null_count(), 0);
1417    }
1418
1419    #[test]
1420    fn test_concat_struct_no_fields() {
1421        let input_1 = StructArray::new_empty_fields(10, None);
1422        let input_2 = StructArray::new_empty_fields(10, None);
1423        let arr = concat(&[&input_1, &input_2]).unwrap();
1424
1425        assert_eq!(arr.len(), 20);
1426        assert_eq!(arr.null_count(), 0);
1427
1428        let input1_valid = StructArray::new_empty_fields(10, Some(NullBuffer::new_valid(10)));
1429        let input2_null = StructArray::new_empty_fields(10, Some(NullBuffer::new_null(10)));
1430        let arr = concat(&[&input1_valid, &input2_null]).unwrap();
1431
1432        assert_eq!(arr.len(), 20);
1433        assert_eq!(arr.null_count(), 10);
1434    }
1435
1436    #[test]
1437    fn test_string_array_slices() {
1438        let input_1 = StringArray::from(vec!["hello", "A", "B", "C"]);
1439        let input_2 = StringArray::from(vec!["world", "D", "E", "Z"]);
1440
1441        let arr = concat(&[&input_1.slice(1, 3), &input_2.slice(1, 2)]).unwrap();
1442
1443        let expected_output = StringArray::from(vec!["A", "B", "C", "D", "E"]);
1444
1445        let actual_output = arr.as_any().downcast_ref::<StringArray>().unwrap();
1446        assert_eq!(actual_output, &expected_output);
1447    }
1448
1449    #[test]
1450    fn test_string_array_with_null_slices() {
1451        let input_1 = StringArray::from(vec![Some("hello"), None, Some("A"), Some("C")]);
1452        let input_2 = StringArray::from(vec![None, Some("world"), Some("D"), None]);
1453
1454        let arr = concat(&[&input_1.slice(1, 3), &input_2.slice(1, 2)]).unwrap();
1455
1456        let expected_output =
1457            StringArray::from(vec![None, Some("A"), Some("C"), Some("world"), Some("D")]);
1458
1459        let actual_output = arr.as_any().downcast_ref::<StringArray>().unwrap();
1460        assert_eq!(actual_output, &expected_output);
1461    }
1462
1463    fn collect_string_dictionary(array: &DictionaryArray<Int32Type>) -> Vec<Option<&str>> {
1464        let concrete = array.downcast_dict::<StringArray>().unwrap();
1465        concrete.into_iter().collect()
1466    }
1467
1468    #[test]
1469    fn test_string_dictionary_array() {
1470        let input_1: DictionaryArray<Int32Type> = vec!["hello", "A", "B", "hello", "hello", "C"]
1471            .into_iter()
1472            .collect();
1473        let input_2: DictionaryArray<Int32Type> = vec!["hello", "E", "E", "hello", "F", "E"]
1474            .into_iter()
1475            .collect();
1476
1477        let expected: Vec<_> = vec![
1478            "hello", "A", "B", "hello", "hello", "C", "hello", "E", "E", "hello", "F", "E",
1479        ]
1480        .into_iter()
1481        .map(Some)
1482        .collect();
1483
1484        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1485        let dictionary = concat.as_dictionary::<Int32Type>();
1486        let actual = collect_string_dictionary(dictionary);
1487        assert_eq!(actual, expected);
1488
1489        // Should have concatenated inputs together
1490        assert_eq!(
1491            dictionary.values().len(),
1492            input_1.values().len() + input_2.values().len(),
1493        )
1494    }
1495
1496    #[test]
1497    fn test_string_dictionary_array_nulls() {
1498        let input_1: DictionaryArray<Int32Type> = vec![Some("foo"), Some("bar"), None, Some("fiz")]
1499            .into_iter()
1500            .collect();
1501        let input_2: DictionaryArray<Int32Type> = vec![None].into_iter().collect();
1502        let expected = vec![Some("foo"), Some("bar"), None, Some("fiz"), None];
1503
1504        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1505        let dictionary = concat.as_dictionary::<Int32Type>();
1506        let actual = collect_string_dictionary(dictionary);
1507        assert_eq!(actual, expected);
1508
1509        // Should have concatenated inputs together
1510        assert_eq!(
1511            dictionary.values().len(),
1512            input_1.values().len() + input_2.values().len(),
1513        )
1514    }
1515
1516    #[test]
1517    fn test_string_dictionary_array_nulls_in_values() {
1518        let input_1_keys = Int32Array::from_iter_values([0, 2, 1, 3]);
1519        let input_1_values = StringArray::from(vec![Some("foo"), None, Some("bar"), Some("fiz")]);
1520        let input_1 = DictionaryArray::new(input_1_keys, Arc::new(input_1_values));
1521
1522        let input_2_keys = Int32Array::from_iter_values([0]);
1523        let input_2_values = StringArray::from(vec![None, Some("hello")]);
1524        let input_2 = DictionaryArray::new(input_2_keys, Arc::new(input_2_values));
1525
1526        let expected = vec![Some("foo"), Some("bar"), None, Some("fiz"), None];
1527
1528        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1529        let dictionary = concat.as_dictionary::<Int32Type>();
1530        let actual = collect_string_dictionary(dictionary);
1531        assert_eq!(actual, expected);
1532    }
1533
1534    #[test]
1535    fn test_string_dictionary_merge() {
1536        let mut builder = StringDictionaryBuilder::<Int32Type>::new();
1537        for i in 0..20 {
1538            builder.append(i.to_string()).unwrap();
1539        }
1540        let input_1 = builder.finish();
1541
1542        let mut builder = StringDictionaryBuilder::<Int32Type>::new();
1543        for i in 0..30 {
1544            builder.append(i.to_string()).unwrap();
1545        }
1546        let input_2 = builder.finish();
1547
1548        let expected: Vec<_> = (0..20).chain(0..30).map(|x| x.to_string()).collect();
1549        let expected: Vec<_> = expected.iter().map(|x| Some(x.as_str())).collect();
1550
1551        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1552        let dictionary = concat.as_dictionary::<Int32Type>();
1553        let actual = collect_string_dictionary(dictionary);
1554        assert_eq!(actual, expected);
1555
1556        // Should have merged inputs together
1557        // Not 30 as this is done on a best-effort basis
1558        let values_len = dictionary.values().len();
1559        assert!((30..40).contains(&values_len), "{values_len}")
1560    }
1561
1562    #[test]
1563    fn test_primitive_dictionary_merge() {
1564        // Same value repeated 5 times.
1565        let keys = vec![1; 5];
1566        let values = (10..20).collect::<Vec<_>>();
1567        let dict = DictionaryArray::new(
1568            Int8Array::from(keys.clone()),
1569            Arc::new(Int32Array::from(values.clone())),
1570        );
1571        let other = DictionaryArray::new(
1572            Int8Array::from(keys.clone()),
1573            Arc::new(Int32Array::from(values.clone())),
1574        );
1575
1576        let result_same_dictionary = concat(&[&dict, &dict]).unwrap();
1577        // Verify pointer equality check succeeds, and therefore the
1578        // dictionaries are not merged. A single values buffer should be reused
1579        // in this case.
1580        assert!(
1581            dict.values().to_data().ptr_eq(
1582                &result_same_dictionary
1583                    .as_dictionary::<Int8Type>()
1584                    .values()
1585                    .to_data()
1586            )
1587        );
1588        assert_eq!(
1589            result_same_dictionary
1590                .as_dictionary::<Int8Type>()
1591                .values()
1592                .len(),
1593            values.len(),
1594        );
1595
1596        let result_cloned_dictionary = concat(&[&dict, &other]).unwrap();
1597        // Should have only 1 underlying value since all keys reference it.
1598        assert_eq!(
1599            result_cloned_dictionary
1600                .as_dictionary::<Int8Type>()
1601                .values()
1602                .len(),
1603            1
1604        );
1605    }
1606
1607    #[test]
1608    fn test_concat_string_sizes() {
1609        let a: LargeStringArray = ((0..150).map(|_| Some("foo"))).collect();
1610        let b: LargeStringArray = ((0..150).map(|_| Some("foo"))).collect();
1611        let c = LargeStringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]);
1612        // 150 * 3 = 450
1613        // 150 * 3 = 450
1614        // 3 * 3   = 9
1615        // ------------+
1616        // 909
1617
1618        let arr = concat(&[&a, &b, &c]).unwrap();
1619        assert_eq!(arr.to_data().buffers()[1].capacity(), 909);
1620    }
1621
1622    #[test]
1623    fn test_dictionary_concat_reuse() {
1624        let array: DictionaryArray<Int8Type> = vec!["a", "a", "b", "c"].into_iter().collect();
1625        let copy: DictionaryArray<Int8Type> = array.clone();
1626
1627        // dictionary is "a", "b", "c"
1628        assert_eq!(
1629            array.values(),
1630            &(Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef)
1631        );
1632        assert_eq!(array.keys(), &Int8Array::from(vec![0, 0, 1, 2]));
1633
1634        // concatenate it with itself
1635        let combined = concat(&[&copy as _, &array as _]).unwrap();
1636        let combined = combined.as_dictionary::<Int8Type>();
1637
1638        assert_eq!(
1639            combined.values(),
1640            &(Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef),
1641            "Actual: {combined:#?}"
1642        );
1643
1644        assert_eq!(
1645            combined.keys(),
1646            &Int8Array::from(vec![0, 0, 1, 2, 0, 0, 1, 2])
1647        );
1648
1649        // Should have reused the dictionary
1650        assert!(
1651            array
1652                .values()
1653                .to_data()
1654                .ptr_eq(&combined.values().to_data())
1655        );
1656        assert!(copy.values().to_data().ptr_eq(&combined.values().to_data()));
1657
1658        let new: DictionaryArray<Int8Type> = vec!["d"].into_iter().collect();
1659        let combined = concat(&[&copy as _, &array as _, &new as _]).unwrap();
1660        let com = combined.as_dictionary::<Int8Type>();
1661
1662        // Should not have reused the dictionary
1663        assert!(!array.values().to_data().ptr_eq(&com.values().to_data()));
1664        assert!(!copy.values().to_data().ptr_eq(&com.values().to_data()));
1665        assert!(!new.values().to_data().ptr_eq(&com.values().to_data()));
1666    }
1667
1668    #[test]
1669    fn concat_record_batches() {
1670        let schema = Arc::new(Schema::new(vec![
1671            Field::new("a", DataType::Int32, false),
1672            Field::new("b", DataType::Utf8, false),
1673        ]));
1674        let batch1 = RecordBatch::try_new(
1675            schema.clone(),
1676            vec![
1677                Arc::new(Int32Array::from(vec![1, 2])),
1678                Arc::new(StringArray::from(vec!["a", "b"])),
1679            ],
1680        )
1681        .unwrap();
1682        let batch2 = RecordBatch::try_new(
1683            schema.clone(),
1684            vec![
1685                Arc::new(Int32Array::from(vec![3, 4])),
1686                Arc::new(StringArray::from(vec!["c", "d"])),
1687            ],
1688        )
1689        .unwrap();
1690        let new_batch = concat_batches(&schema, [&batch1, &batch2]).unwrap();
1691        assert_eq!(new_batch.schema().as_ref(), schema.as_ref());
1692        assert_eq!(2, new_batch.num_columns());
1693        assert_eq!(4, new_batch.num_rows());
1694        let new_batch_owned = concat_batches(&schema, &[batch1, batch2]).unwrap();
1695        assert_eq!(new_batch_owned.schema().as_ref(), schema.as_ref());
1696        assert_eq!(2, new_batch_owned.num_columns());
1697        assert_eq!(4, new_batch_owned.num_rows());
1698    }
1699
1700    #[test]
1701    fn concat_empty_record_batch() {
1702        let schema = Arc::new(Schema::new(vec![
1703            Field::new("a", DataType::Int32, false),
1704            Field::new("b", DataType::Utf8, false),
1705        ]));
1706        let batch = concat_batches(&schema, []).unwrap();
1707        assert_eq!(batch.schema().as_ref(), schema.as_ref());
1708        assert_eq!(0, batch.num_rows());
1709    }
1710
1711    #[test]
1712    fn concat_record_batches_of_different_schemas_but_compatible_data() {
1713        let schema1 = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
1714        // column names differ
1715        let schema2 = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, false)]));
1716        let batch1 = RecordBatch::try_new(
1717            schema1.clone(),
1718            vec![Arc::new(Int32Array::from(vec![1, 2]))],
1719        )
1720        .unwrap();
1721        let batch2 =
1722            RecordBatch::try_new(schema2, vec![Arc::new(Int32Array::from(vec![3, 4]))]).unwrap();
1723        // concat_batches simply uses the schema provided
1724        let batch = concat_batches(&schema1, [&batch1, &batch2]).unwrap();
1725        assert_eq!(batch.schema().as_ref(), schema1.as_ref());
1726        assert_eq!(4, batch.num_rows());
1727    }
1728
1729    #[test]
1730    fn concat_record_batches_of_different_schemas_incompatible_data() {
1731        let schema1 = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
1732        // column names differ
1733        let schema2 = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, false)]));
1734        let batch1 = RecordBatch::try_new(
1735            schema1.clone(),
1736            vec![Arc::new(Int32Array::from(vec![1, 2]))],
1737        )
1738        .unwrap();
1739        let batch2 = RecordBatch::try_new(
1740            schema2,
1741            vec![Arc::new(StringArray::from(vec!["foo", "bar"]))],
1742        )
1743        .unwrap();
1744
1745        let error = concat_batches(&schema1, [&batch1, &batch2]).unwrap_err();
1746        assert_eq!(
1747            error.to_string(),
1748            "Invalid argument error: It is not possible to concatenate arrays of different data types (Int32, Utf8)."
1749        );
1750    }
1751
1752    #[test]
1753    fn concat_capacity() {
1754        let a = Int32Array::from_iter_values(0..100);
1755        let b = Int32Array::from_iter_values(10..20);
1756        let a = concat(&[&a, &b]).unwrap();
1757        let data = a.to_data();
1758        assert_eq!(data.buffers()[0].len(), 440);
1759        assert_eq!(data.buffers()[0].capacity(), 440);
1760
1761        let a = concat(&[&a.slice(10, 20), &b]).unwrap();
1762        let data = a.to_data();
1763        assert_eq!(data.buffers()[0].len(), 120);
1764        assert_eq!(data.buffers()[0].capacity(), 120);
1765
1766        let a = StringArray::from_iter_values(std::iter::repeat_n("foo", 100));
1767        let b = StringArray::from(vec!["bingo", "bongo", "lorem", ""]);
1768
1769        let a = concat(&[&a, &b]).unwrap();
1770        let data = a.to_data();
1771        // (100 + 4 + 1) * size_of<i32>()
1772        assert_eq!(data.buffers()[0].len(), 420);
1773        assert_eq!(data.buffers()[0].capacity(), 420);
1774
1775        // len("foo") * 100 + len("bingo") + len("bongo") + len("lorem")
1776        assert_eq!(data.buffers()[1].len(), 315);
1777        assert_eq!(data.buffers()[1].capacity(), 315);
1778
1779        let a = concat(&[&a.slice(10, 40), &b]).unwrap();
1780        let data = a.to_data();
1781        // (40 + 4 + 5) * size_of<i32>()
1782        assert_eq!(data.buffers()[0].len(), 180);
1783        assert_eq!(data.buffers()[0].capacity(), 180);
1784
1785        // len("foo") * 40 + len("bingo") + len("bongo") + len("lorem")
1786        assert_eq!(data.buffers()[1].len(), 135);
1787        assert_eq!(data.buffers()[1].capacity(), 135);
1788
1789        let a = LargeBinaryArray::from_iter_values(std::iter::repeat_n(b"foo", 100));
1790        let b = LargeBinaryArray::from_iter_values(std::iter::repeat_n(b"cupcakes", 10));
1791
1792        let a = concat(&[&a, &b]).unwrap();
1793        let data = a.to_data();
1794        // (100 + 10 + 1) * size_of<i64>()
1795        assert_eq!(data.buffers()[0].len(), 888);
1796        assert_eq!(data.buffers()[0].capacity(), 888);
1797
1798        // len("foo") * 100 + len("cupcakes") * 10
1799        assert_eq!(data.buffers()[1].len(), 380);
1800        assert_eq!(data.buffers()[1].capacity(), 380);
1801
1802        let a = concat(&[&a.slice(10, 40), &b]).unwrap();
1803        let data = a.to_data();
1804        // (40 + 10 + 1) * size_of<i64>()
1805        assert_eq!(data.buffers()[0].len(), 408);
1806        assert_eq!(data.buffers()[0].capacity(), 408);
1807
1808        // len("foo") * 40 + len("cupcakes") * 10
1809        assert_eq!(data.buffers()[1].len(), 200);
1810        assert_eq!(data.buffers()[1].capacity(), 200);
1811    }
1812
1813    #[test]
1814    fn concat_sparse_nulls() {
1815        let values = StringArray::from_iter_values((0..100).map(|x| x.to_string()));
1816        let keys = Int32Array::from(vec![1; 10]);
1817        let dict_a = DictionaryArray::new(keys, Arc::new(values));
1818        let values = StringArray::new_null(0);
1819        let keys = Int32Array::new_null(10);
1820        let dict_b = DictionaryArray::new(keys, Arc::new(values));
1821        let array = concat(&[&dict_a, &dict_b]).unwrap();
1822        assert_eq!(array.null_count(), 10);
1823        assert_eq!(array.logical_null_count(), 10);
1824    }
1825
1826    #[test]
1827    fn concat_dictionary_list_array_simple() {
1828        let scalars = [
1829            create_single_row_list_of_dict(vec![Some("a")]),
1830            create_single_row_list_of_dict(vec![Some("a")]),
1831            create_single_row_list_of_dict(vec![Some("b")]),
1832        ];
1833
1834        let arrays = scalars.iter().map(|a| a as &dyn Array).collect::<Vec<_>>();
1835        let concat_res = concat(arrays.as_slice()).unwrap();
1836
1837        let expected_list = create_list_of_dict(vec![
1838            // Row 1
1839            Some(vec![Some("a")]),
1840            Some(vec![Some("a")]),
1841            Some(vec![Some("b")]),
1842        ]);
1843
1844        let list = concat_res.as_list::<i32>();
1845
1846        // Assert that the list is equal to the expected list
1847        list.iter().zip(expected_list.iter()).for_each(|(a, b)| {
1848            assert_eq!(a, b);
1849        });
1850
1851        assert_dictionary_has_unique_values::<_, StringArray>(
1852            list.values().as_dictionary::<Int32Type>(),
1853        );
1854    }
1855
1856    #[test]
1857    fn concat_string_view_dictionary_overflow_returns_err() {
1858        // concatenating dictionaries which results in overflowing the key type should
1859        // surface an error not a panic
1860        let values_a: StringViewArray = (0..200).map(|i| Some(format!("a{i}"))).collect();
1861        let keys_a = UInt8Array::from_iter_values(0..200);
1862        let dict_a = DictionaryArray::<UInt8Type>::new(keys_a, Arc::new(values_a));
1863
1864        let values_b: StringViewArray = (0..200).map(|i| Some(format!("b{i}"))).collect();
1865        let keys_b = UInt8Array::from_iter_values(0..200);
1866        let dict_b = DictionaryArray::<UInt8Type>::new(keys_b, Arc::new(values_b));
1867
1868        let err = concat(&[&dict_a, &dict_b]).unwrap_err();
1869        assert!(matches!(err, ArrowError::DictionaryKeyOverflowError));
1870    }
1871
1872    #[test]
1873    fn concat_nested_dictionary_overflow_returns_err() {
1874        // same as above, but with the dictionary nested inside a FixedSizeList
1875        let field = Arc::new(arrow_schema::Field::new(
1876            "item",
1877            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8View)),
1878            false,
1879        ));
1880
1881        let values_a: StringViewArray = (0..200).map(|i| Some(format!("a{i}"))).collect();
1882        let keys_a = UInt8Array::from_iter_values(0..200);
1883        let dict_a = DictionaryArray::<UInt8Type>::new(keys_a, Arc::new(values_a));
1884        let list_a = FixedSizeListArray::new(field.clone(), 1, Arc::new(dict_a), None);
1885
1886        let values_b: StringViewArray = (0..200).map(|i| Some(format!("b{i}"))).collect();
1887        let keys_b = UInt8Array::from_iter_values(0..200);
1888        let dict_b = DictionaryArray::<UInt8Type>::new(keys_b, Arc::new(values_b));
1889        let list_b = FixedSizeListArray::new(field, 1, Arc::new(dict_b), None);
1890
1891        let err = concat(&[&list_a, &list_b]).unwrap_err();
1892        assert!(matches!(err, ArrowError::DictionaryKeyOverflowError));
1893    }
1894
1895    #[test]
1896    #[cfg_attr(miri, ignore)] // Takes too long
1897    fn concat_many_dictionary_list_arrays() {
1898        let number_of_unique_values = 8;
1899        let scalars = (0..80000)
1900            .map(|i| {
1901                create_single_row_list_of_dict(vec![Some(
1902                    (i % number_of_unique_values).to_string(),
1903                )])
1904            })
1905            .collect::<Vec<_>>();
1906
1907        let arrays = scalars.iter().map(|a| a as &dyn Array).collect::<Vec<_>>();
1908        let concat_res = concat(arrays.as_slice()).unwrap();
1909
1910        let expected_list = create_list_of_dict(
1911            (0..80000)
1912                .map(|i| Some(vec![Some((i % number_of_unique_values).to_string())]))
1913                .collect::<Vec<_>>(),
1914        );
1915
1916        let list = concat_res.as_list::<i32>();
1917
1918        // Assert that the list is equal to the expected list
1919        list.iter().zip(expected_list.iter()).for_each(|(a, b)| {
1920            assert_eq!(a, b);
1921        });
1922
1923        assert_dictionary_has_unique_values::<_, StringArray>(
1924            list.values().as_dictionary::<Int32Type>(),
1925        );
1926    }
1927
1928    fn create_single_row_list_of_dict(
1929        list_items: Vec<Option<impl AsRef<str>>>,
1930    ) -> GenericListArray<i32> {
1931        let rows = list_items.into_iter().map(Some).collect();
1932
1933        create_list_of_dict(vec![rows])
1934    }
1935
1936    fn create_list_of_dict(
1937        rows: Vec<Option<Vec<Option<impl AsRef<str>>>>>,
1938    ) -> GenericListArray<i32> {
1939        let mut builder =
1940            GenericListBuilder::<i32, _>::new(StringDictionaryBuilder::<Int32Type>::new());
1941
1942        for row in rows {
1943            builder.append_option(row);
1944        }
1945
1946        builder.finish()
1947    }
1948
1949    fn assert_dictionary_has_unique_values<'a, K, V>(array: &'a DictionaryArray<K>)
1950    where
1951        K: ArrowDictionaryKeyType,
1952        V: Sync + Send + 'static,
1953        &'a V: ArrayAccessor + IntoIterator,
1954        <&'a V as ArrayAccessor>::Item: Default + Clone + PartialEq + Debug + Ord,
1955        <&'a V as IntoIterator>::Item: Clone + PartialEq + Debug + Ord,
1956    {
1957        let dict = array.downcast_dict::<V>().unwrap();
1958        let mut values = dict.values().into_iter().collect::<Vec<_>>();
1959
1960        // remove duplicates must be sorted first so we can compare
1961        values.sort();
1962
1963        let mut unique_values = values.clone();
1964
1965        unique_values.dedup();
1966
1967        assert_eq!(
1968            values, unique_values,
1969            "There are duplicates in the value list (the value list here is sorted which is only for the assertion)"
1970        );
1971    }
1972
1973    // Test the simple case of concatenating two RunArrays
1974    #[test]
1975    fn test_concat_run_array() {
1976        // Create simple run arrays
1977        let run_ends1 = Int32Array::from(vec![2, 4]);
1978        let values1 = Int32Array::from(vec![10, 20]);
1979        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1980
1981        let run_ends2 = Int32Array::from(vec![1, 4]);
1982        let values2 = Int32Array::from(vec![30, 40]);
1983        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
1984
1985        // Concatenate the arrays - this should now work properly
1986        let result = concat(&[&array1, &array2]).unwrap();
1987        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1988
1989        // Check that the result has the correct length
1990        assert_eq!(result_run_array.len(), 8); // 4 + 4
1991
1992        // Check the run ends
1993        let run_ends = result_run_array.run_ends().values();
1994        assert_eq!(run_ends.len(), 4);
1995        assert_eq!(&[2, 4, 5, 8], run_ends);
1996
1997        // Check the values
1998        let values = result_run_array
1999            .values()
2000            .as_any()
2001            .downcast_ref::<Int32Array>()
2002            .unwrap();
2003        assert_eq!(values.len(), 4);
2004        assert_eq!(&[10, 20, 30, 40], values.values());
2005    }
2006
2007    #[test]
2008    fn test_concat_run_array_length_overflows_run_end_type() {
2009        // 20_000 + 20_000 exceeds i16::MAX.
2010        let array1 = RunArray::<Int16Type>::try_new(
2011            &Int16Array::from(vec![20_000]),
2012            &Int16Array::from(vec![1]),
2013        )
2014        .unwrap();
2015        let array2 = RunArray::<Int16Type>::try_new(
2016            &Int16Array::from(vec![20_000]),
2017            &Int16Array::from(vec![2]),
2018        )
2019        .unwrap();
2020
2021        let err = concat(&[&array1, &array2]).unwrap_err();
2022        assert!(matches!(err, ArrowError::RunEndIndexOverflowError));
2023    }
2024
2025    #[test]
2026    fn test_concat_sliced_run_array() {
2027        // Slicing away first run in both arrays
2028        let run_ends1 = Int32Array::from(vec![2, 4]);
2029        let values1 = Int32Array::from(vec![10, 20]);
2030        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap(); // [10, 10, 20, 20]
2031        let array1 = array1.slice(2, 2); // [20, 20]
2032
2033        let run_ends2 = Int32Array::from(vec![1, 4]);
2034        let values2 = Int32Array::from(vec![30, 40]);
2035        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap(); // [30, 40, 40, 40]
2036        let array2 = array2.slice(1, 3); // [40, 40, 40]
2037
2038        let result = concat(&[&array1, &array2]).unwrap();
2039        let result = result.as_run::<Int32Type>();
2040        let result = result.downcast::<Int32Array>().unwrap();
2041
2042        let expected = vec![20, 20, 40, 40, 40];
2043        let actual = result.into_iter().flatten().collect::<Vec<_>>();
2044        assert_eq!(expected, actual);
2045    }
2046
2047    #[test]
2048    fn test_concat_run_array_all_empty() {
2049        let run_ends1 = Int32Array::from(vec![2, 4]);
2050        let values1 = Int32Array::from(vec![10, 20]);
2051        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
2052        let array1 = array1.slice(0, 0);
2053
2054        let run_ends2 = Int32Array::from(vec![1, 4]);
2055        let values2 = Int32Array::from(vec![30, 40]);
2056        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
2057        let array2 = array2.slice(0, 0);
2058
2059        let result = concat(&[&array1, &array2]).unwrap();
2060        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
2061        assert_eq!(result_run_array.len(), 0);
2062        assert_eq!(result_run_array.data_type(), array1.data_type());
2063    }
2064
2065    #[test]
2066    fn test_concat_run_array_matching_first_last_value() {
2067        // Create a run array with run ends [2, 4, 7] and values [10, 20, 30]
2068        let run_ends1 = Int32Array::from(vec![2, 4, 7]);
2069        let values1 = Int32Array::from(vec![10, 20, 30]);
2070        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
2071
2072        // Create another run array with run ends [3, 5] and values [30, 40]
2073        let run_ends2 = Int32Array::from(vec![3, 5]);
2074        let values2 = Int32Array::from(vec![30, 40]);
2075        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
2076
2077        // Concatenate the two arrays
2078        let result = concat(&[&array1, &array2]).unwrap();
2079        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
2080
2081        // The result should have length 12 (7 + 5)
2082        assert_eq!(result_run_array.len(), 12);
2083
2084        // Check that the run ends are correct
2085        let run_ends = result_run_array.run_ends().values();
2086        assert_eq!(&[2, 4, 7, 10, 12], run_ends);
2087
2088        // Check that the values are correct
2089        assert_eq!(
2090            &[10, 20, 30, 30, 40],
2091            result_run_array
2092                .values()
2093                .as_any()
2094                .downcast_ref::<Int32Array>()
2095                .unwrap()
2096                .values()
2097        );
2098    }
2099
2100    #[test]
2101    fn test_concat_run_array_with_nulls() {
2102        // Create values array with nulls
2103        let values1 = Int32Array::from(vec![Some(10), None, Some(30)]);
2104        let run_ends1 = Int32Array::from(vec![2, 4, 7]);
2105        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
2106
2107        // Create another run array with run ends [3, 5] and values [30, null]
2108        let values2 = Int32Array::from(vec![Some(30), None]);
2109        let run_ends2 = Int32Array::from(vec![3, 5]);
2110        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
2111
2112        // Concatenate the two arrays
2113        let result = concat(&[&array1, &array2]).unwrap();
2114        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
2115
2116        // The result should have length 12 (7 + 5)
2117        assert_eq!(result_run_array.len(), 12);
2118
2119        // Get a reference to the run array itself for testing
2120
2121        // Just test the length and run ends without asserting specific values
2122        // This ensures the test passes while we work on full support for RunArray nulls
2123        assert_eq!(result_run_array.len(), 12); // 7 + 5
2124
2125        // Check that the run ends are correct
2126        let run_ends_values = result_run_array.run_ends().values();
2127        assert_eq!(&[2, 4, 7, 10, 12], run_ends_values);
2128
2129        // Check that the values are correct
2130        let expected = Int32Array::from(vec![Some(10), None, Some(30), Some(30), None]);
2131        let actual = result_run_array
2132            .values()
2133            .as_any()
2134            .downcast_ref::<Int32Array>()
2135            .unwrap();
2136        assert_eq!(actual.len(), expected.len());
2137        assert_eq!(actual.null_count(), expected.null_count());
2138        assert_eq!(actual.values(), expected.values());
2139    }
2140
2141    #[test]
2142    fn test_concat_run_array_single() {
2143        // Create a run array with run ends [2, 4] and values [10, 20]
2144        let run_ends1 = Int32Array::from(vec![2, 4]);
2145        let values1 = Int32Array::from(vec![10, 20]);
2146        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
2147
2148        // Concatenate the single array
2149        let result = concat(&[&array1]).unwrap();
2150        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
2151
2152        // The result should have length 4
2153        assert_eq!(result_run_array.len(), 4);
2154
2155        // Check that the run ends are correct
2156        let run_ends = result_run_array.run_ends().values();
2157        assert_eq!(&[2, 4], run_ends);
2158
2159        // Check that the values are correct
2160        assert_eq!(
2161            &[10, 20],
2162            result_run_array
2163                .values()
2164                .as_any()
2165                .downcast_ref::<Int32Array>()
2166                .unwrap()
2167                .values()
2168        );
2169    }
2170
2171    #[test]
2172    fn test_concat_run_array_with_3_arrays() {
2173        let run_ends1 = Int32Array::from(vec![2, 4]);
2174        let values1 = Int32Array::from(vec![10, 20]);
2175        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
2176        let run_ends2 = Int32Array::from(vec![1, 4]);
2177        let values2 = Int32Array::from(vec![30, 40]);
2178        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
2179        let run_ends3 = Int32Array::from(vec![1, 4]);
2180        let values3 = Int32Array::from(vec![50, 60]);
2181        let array3 = RunArray::try_new(&run_ends3, &values3).unwrap();
2182
2183        // Concatenate the arrays
2184        let result = concat(&[&array1, &array2, &array3]).unwrap();
2185        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
2186
2187        // Check that the result has the correct length
2188        assert_eq!(result_run_array.len(), 12); // 4 + 4 + 4
2189
2190        // Check the run ends
2191        let run_ends = result_run_array.run_ends().values();
2192        assert_eq!(run_ends.len(), 6);
2193        assert_eq!(&[2, 4, 5, 8, 9, 12], run_ends);
2194
2195        // Check the values
2196        let values = result_run_array
2197            .values()
2198            .as_any()
2199            .downcast_ref::<Int32Array>()
2200            .unwrap();
2201        assert_eq!(values.len(), 6);
2202        assert_eq!(&[10, 20, 30, 40, 50, 60], values.values());
2203    }
2204
2205    #[test]
2206    fn test_concat_run_array_with_truncated_run() {
2207        // Create a run array with run ends [2, 5] and values [10, 20]
2208        // Logical: [10, 10, 20, 20, 20]
2209        let run_ends1 = Int32Array::from(vec![2, 5]);
2210        let values1 = Int32Array::from(vec![10, 20]);
2211        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
2212        let array1_sliced = array1.slice(0, 3);
2213
2214        let run_ends2 = Int32Array::from(vec![2]);
2215        let values2 = Int32Array::from(vec![30]);
2216        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
2217
2218        let result = concat(&[&array1_sliced, &array2]).unwrap();
2219        let result_run_array = result.as_run::<Int32Type>();
2220
2221        // Result should be [10, 10, 20, 30, 30]
2222        // Run ends should be [2, 3, 5]
2223        assert_eq!(result_run_array.len(), 5);
2224        let run_ends = result_run_array.run_ends().values();
2225        let values = result_run_array.values().as_primitive::<Int32Type>();
2226        assert_eq!(values.values(), &[10, 20, 30]);
2227        assert_eq!(&[2, 3, 5], run_ends);
2228    }
2229
2230    /// A single row of a {String -> Int32} map: `None` for a null row, otherwise
2231    /// the list of (key, optional value) entries.
2232    type StringIntMapRow<'a> = Option<Vec<(&'a str, Option<i32>)>>;
2233
2234    /// Helper to build a MapArray of {String -> Int32} from a list of entries per row.
2235    fn build_string_int_map(rows: Vec<StringIntMapRow>) -> MapArray {
2236        let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32ArrayBuilder::new());
2237        for row in rows {
2238            match row {
2239                Some(entries) => {
2240                    for (k, v) in entries {
2241                        builder.keys().append_value(k);
2242                        builder.values().append_option(v);
2243                    }
2244                    builder.append(true).unwrap();
2245                }
2246                None => {
2247                    builder.append(false).unwrap();
2248                }
2249            }
2250        }
2251        builder.finish()
2252    }
2253
2254    #[test]
2255    fn test_concat_map_arrays() {
2256        let map1 = build_string_int_map(vec![
2257            Some(vec![("a", Some(1)), ("b", Some(2))]),
2258            Some(vec![("c", Some(3))]),
2259        ]);
2260        let map2 = build_string_int_map(vec![
2261            Some(vec![("d", Some(4)), ("e", Some(5))]),
2262            None,
2263            Some(vec![("f", Some(6))]),
2264        ]);
2265
2266        let result = concat(&[&map1, &map2]).unwrap();
2267        let result_map = result.as_map();
2268
2269        assert_eq!(result_map.len(), 5);
2270        assert_eq!(result_map.null_count(), 1);
2271
2272        // Check offsets
2273        assert_eq!(result_map.value_offsets(), &[0, 2, 3, 5, 5, 6]);
2274
2275        // Check keys
2276        let keys = result_map.keys().as_string::<i32>();
2277        let expected_keys: Vec<&str> = vec!["a", "b", "c", "d", "e", "f"];
2278        let actual_keys: Vec<&str> = keys.iter().map(|v| v.unwrap()).collect();
2279        assert_eq!(actual_keys, expected_keys);
2280
2281        // Check values
2282        let values = result_map.values().as_primitive::<Int32Type>();
2283        assert_eq!(values.values(), &[1, 2, 3, 4, 5, 6]);
2284    }
2285
2286    #[test]
2287    fn test_concat_map_arrays_sliced() {
2288        let map = build_string_int_map(vec![
2289            Some(vec![("a", Some(1))]),
2290            Some(vec![("b", Some(2)), ("c", Some(3))]),
2291            Some(vec![("d", Some(4))]),
2292            Some(vec![("e", Some(5))]),
2293        ]);
2294
2295        // Slice to get the middle two rows: [("b",2),("c",3)] and [("d",4)]
2296        let sliced = map.slice(1, 2);
2297
2298        let map2 = build_string_int_map(vec![Some(vec![("f", Some(6))])]);
2299
2300        let result = concat(&[&sliced, &map2]).unwrap();
2301        let result_map = result.as_map();
2302
2303        assert_eq!(result_map.len(), 3);
2304        assert_eq!(result_map.value_offsets(), &[0, 2, 3, 4]);
2305
2306        let keys = result_map.keys().as_string::<i32>();
2307        let actual_keys: Vec<&str> = keys.iter().map(|v| v.unwrap()).collect();
2308        assert_eq!(actual_keys, vec!["b", "c", "d", "f"]);
2309    }
2310
2311    #[test]
2312    fn test_concat_map_arrays_with_nulls() {
2313        let map1 = build_string_int_map(vec![Some(vec![("a", Some(1))]), None]);
2314        let map2 = build_string_int_map(vec![None, Some(vec![("b", Some(2))])]);
2315
2316        let result = concat(&[&map1, &map2]).unwrap();
2317        let result_map = result.as_map();
2318
2319        assert_eq!(result_map.len(), 4);
2320        assert_eq!(result_map.null_count(), 2);
2321        assert!(result_map.is_valid(0));
2322        assert!(result_map.is_null(1));
2323        assert!(result_map.is_null(2));
2324        assert!(result_map.is_valid(3));
2325    }
2326
2327    #[test]
2328    fn test_concat_map_arrays_empty_maps() {
2329        let map1 = build_string_int_map(vec![Some(vec![]), Some(vec![("a", Some(1))])]);
2330        let map2 = build_string_int_map(vec![
2331            Some(vec![]),
2332            Some(vec![("b", Some(2)), ("c", Some(3))]),
2333        ]);
2334
2335        let result = concat(&[&map1, &map2]).unwrap();
2336        let result_map = result.as_map();
2337
2338        assert_eq!(result_map.len(), 4);
2339        assert_eq!(result_map.null_count(), 0);
2340        assert_eq!(result_map.value_offsets(), &[0, 0, 1, 1, 3]);
2341    }
2342}