Skip to main content

arrow_select/
concat.rs

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