Skip to main content

arrow_select/
concat.rs

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