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.iter() {
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().unwrap().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().unwrap().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().unwrap().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().unwrap().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.iter() {
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 (item_capacity, bytes_capacity) = match binary_capacity::<T>(arrays) {
357        Capacities::Binary(item_capacity, Some(bytes_capacity)) => (item_capacity, bytes_capacity),
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    // The run ends need to be adjusted by the sum of the lengths of the previous arrays.
428    let needed_run_end_adjustments = std::iter::once(R::default_value())
429        .chain(
430            run_arrays
431                .iter()
432                .scan(R::default_value(), |acc, run_array| {
433                    *acc = *acc + R::Native::from_usize(run_array.len()).unwrap();
434                    Some(*acc)
435                }),
436        )
437        .collect::<Vec<_>>();
438
439    // This works out nicely to be the total (logical) length of the resulting array.
440    let total_len = needed_run_end_adjustments.last().unwrap().as_usize();
441
442    let run_ends_array =
443        PrimitiveArray::<R>::from_iter_values(run_arrays.iter().enumerate().flat_map(
444            move |(i, run_array)| {
445                let adjustment = needed_run_end_adjustments[i];
446                run_array
447                    .run_ends()
448                    .sliced_values()
449                    .map(move |run_end| run_end + adjustment)
450            },
451        ));
452
453    let values_slices: Vec<ArrayRef> = run_arrays
454        .iter()
455        .map(|run_array| run_array.values_slice())
456        .collect();
457
458    let all_values = concat(&values_slices.iter().map(|x| x.as_ref()).collect::<Vec<_>>())?;
459
460    let builder = ArrayDataBuilder::new(run_arrays[0].data_type().clone())
461        .len(total_len)
462        .child_data(vec![run_ends_array.into_data(), all_values.into_data()]);
463
464    // `build_unchecked` is used to avoid recursive validation of child arrays.
465    let array_data = unsafe { builder.build_unchecked() };
466    array_data.validate_data()?;
467
468    Ok(Arc::<RunArray<R>>::new(array_data.into()))
469}
470
471macro_rules! dict_helper {
472    ($t:ty, $arrays:expr) => {
473        return concat_dictionaries::<$t>($arrays)
474    };
475}
476
477macro_rules! primitive_concat {
478    ($t:ty, $arrays:expr) => {
479        return concat_primitives::<$t>($arrays)
480    };
481}
482
483fn get_capacity(arrays: &[&dyn Array], data_type: &DataType) -> Capacities {
484    match data_type {
485        DataType::Utf8 => binary_capacity::<Utf8Type>(arrays),
486        DataType::LargeUtf8 => binary_capacity::<LargeUtf8Type>(arrays),
487        DataType::Binary => binary_capacity::<BinaryType>(arrays),
488        DataType::LargeBinary => binary_capacity::<LargeBinaryType>(arrays),
489        DataType::FixedSizeList(_, _) => fixed_size_list_capacity(arrays, data_type),
490        _ => Capacities::Array(arrays.iter().map(|a| a.len()).sum()),
491    }
492}
493
494/// Concatenate multiple [Array] of the same type into a single [ArrayRef].
495pub fn concat(arrays: &[&dyn Array]) -> Result<ArrayRef, ArrowError> {
496    if arrays.is_empty() {
497        return Err(ArrowError::ComputeError(
498            "concat requires input of at least one array".to_string(),
499        ));
500    } else if arrays.len() == 1 {
501        let array = arrays[0];
502        return Ok(array.slice(0, array.len()));
503    }
504
505    let d = arrays[0].data_type();
506    if arrays.iter().skip(1).any(|array| array.data_type() != d) {
507        // Create error message with up to 10 unique data types in the order they appear
508        let error_message = {
509            // 10 max unique data types to print and another 1 to know if there are more
510            let mut unique_data_types = HashSet::with_capacity(11);
511
512            let mut error_message =
513                format!("It is not possible to concatenate arrays of different data types ({d}");
514            unique_data_types.insert(d);
515
516            for array in arrays {
517                let is_unique = unique_data_types.insert(array.data_type());
518
519                if unique_data_types.len() == 11 {
520                    error_message.push_str(", ...");
521                    break;
522                }
523
524                if is_unique {
525                    error_message.push_str(", ");
526                    error_message.push_str(&array.data_type().to_string());
527                }
528            }
529
530            error_message.push_str(").");
531
532            error_message
533        };
534
535        return Err(ArrowError::InvalidArgumentError(error_message));
536    }
537
538    downcast_primitive! {
539        d => (primitive_concat, arrays),
540        DataType::Boolean => concat_boolean(arrays),
541        DataType::Dictionary(k, _) => {
542            downcast_integer! {
543                k.as_ref() => (dict_helper, arrays),
544                _ => unreachable!("illegal dictionary key type {k}")
545            }
546        }
547        DataType::List(field) => concat_lists::<i32>(arrays, field),
548        DataType::LargeList(field) => concat_lists::<i64>(arrays, field),
549        DataType::ListView(field) => concat_list_view::<i32>(arrays, field),
550        DataType::LargeListView(field) => concat_list_view::<i64>(arrays, field),
551        DataType::Map(field, ordered) => concat_maps(arrays, field, *ordered),
552        DataType::Struct(fields) => concat_structs(arrays, fields),
553        DataType::Utf8 => concat_bytes::<Utf8Type>(arrays),
554        DataType::LargeUtf8 => concat_bytes::<LargeUtf8Type>(arrays),
555        DataType::Binary => concat_bytes::<BinaryType>(arrays),
556        DataType::LargeBinary => concat_bytes::<LargeBinaryType>(arrays),
557        DataType::RunEndEncoded(r, _) => {
558            // Handle RunEndEncoded arrays with special concat function
559            // We need to downcast based on the run end type
560            match r.data_type() {
561                DataType::Int16 => concat_run_arrays::<Int16Type>(arrays),
562                DataType::Int32 => concat_run_arrays::<Int32Type>(arrays),
563                DataType::Int64 => concat_run_arrays::<Int64Type>(arrays),
564                _ => unreachable!("Unsupported run end index type: {r:?}"),
565            }
566        }
567        DataType::Utf8View => concat_byte_view::<StringViewType>(arrays),
568        DataType::BinaryView => concat_byte_view::<BinaryViewType>(arrays),
569        _ => {
570            let capacity = get_capacity(arrays, d);
571            concat_fallback(arrays, capacity)
572        }
573    }
574}
575
576/// Concatenates arrays using MutableArrayData
577///
578/// This will naively concatenate dictionaries
579fn concat_fallback(arrays: &[&dyn Array], capacity: Capacities) -> Result<ArrayRef, ArrowError> {
580    let array_data: Vec<_> = arrays.iter().map(|a| a.to_data()).collect::<Vec<_>>();
581    let array_data = array_data.iter().collect();
582    let mut mutable = MutableArrayData::with_capacities(array_data, false, capacity);
583
584    for (i, a) in arrays.iter().enumerate() {
585        mutable.try_extend(i, 0, a.len())?
586    }
587
588    Ok(make_array(mutable.freeze()))
589}
590
591/// Concatenates `batches` together into a single [`RecordBatch`].
592///
593/// The output batch has the specified `schemas`; The schema of the
594/// input are ignored.
595///
596/// # Notes
597///
598/// - Callers should budget for peak memory use to approach 2x the input
599///   size, as the input batches and output arrays co-exist during construction.
600/// - Arrays with `i32` offsets, such as `StringArray` and `BinaryArray`, only
601///   support up to ~2GiB of payloads. Concatenating large arrays of these types
602///   can cause offset overflows.
603///
604/// # Errors
605///
606/// Returns an error if the types of underlying arrays are different.
607pub fn concat_batches<'a>(
608    schema: &SchemaRef,
609    input_batches: impl IntoIterator<Item = &'a RecordBatch>,
610) -> Result<RecordBatch, ArrowError> {
611    // When schema is empty, sum the number of the rows of all batches
612    if schema.fields().is_empty() {
613        let num_rows: usize = input_batches.into_iter().map(RecordBatch::num_rows).sum();
614        let mut options = RecordBatchOptions::default();
615        options.row_count = Some(num_rows);
616        return RecordBatch::try_new_with_options(schema.clone(), vec![], &options);
617    }
618
619    let batches: Vec<&RecordBatch> = input_batches.into_iter().collect();
620    if batches.is_empty() {
621        return Ok(RecordBatch::new_empty(schema.clone()));
622    }
623    let field_num = schema.fields().len();
624    let mut arrays = Vec::with_capacity(field_num);
625    for i in 0..field_num {
626        let array = concat(
627            &batches
628                .iter()
629                .map(|batch| batch.column(i).as_ref())
630                .collect::<Vec<_>>(),
631        )?;
632        arrays.push(array);
633    }
634    RecordBatch::try_new(schema.clone(), arrays)
635}
636
637#[cfg(test)]
638mod tests {
639    use super::*;
640    use arrow_array::builder::{
641        GenericListBuilder, Int32Builder as Int32ArrayBuilder, Int64Builder, ListViewBuilder,
642        MapBuilder, StringBuilder, StringDictionaryBuilder,
643    };
644    use arrow_schema::{Field, Schema};
645    use std::fmt::Debug;
646
647    #[test]
648    fn test_dict_overflow_9366() {
649        use arrow_schema::DataType;
650
651        let schema = Arc::new(Schema::new(vec![Field::new(
652            "a",
653            DataType::Dictionary(
654                Box::new(DataType::UInt8),
655                Box::new(DataType::FixedSizeBinary(8)),
656            ),
657            false,
658        )]));
659        let make = |vals: std::ops::Range<u64>| {
660            let dict = FixedSizeBinaryArray::try_from_iter(vals.map(|i| i.to_le_bytes())).unwrap();
661            let keys = UInt8Array::from_iter_values(0..128);
662            let arr = DictionaryArray::try_new(keys, Arc::new(dict)).unwrap();
663            RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap()
664        };
665        // 256 distinct values fit in u8 keys (0..=255): concat must succeed.
666        let out = concat_batches(&schema, &[make(0..128), make(128..256)]).unwrap();
667        assert_eq!(out.num_rows(), 256);
668        let dict = out.column(0).as_dictionary::<UInt8Type>();
669        assert_eq!(dict.values().len(), 256);
670    }
671
672    #[test]
673    fn test_dict_overflow_i8_9366() {
674        use arrow_schema::DataType;
675
676        // Same boundary for a signed key type: i8 holds 128 keys (0..=127).
677        let schema = Arc::new(Schema::new(vec![Field::new(
678            "a",
679            DataType::Dictionary(
680                Box::new(DataType::Int8),
681                Box::new(DataType::FixedSizeBinary(8)),
682            ),
683            false,
684        )]));
685        let make = |vals: std::ops::Range<u64>| {
686            let dict = FixedSizeBinaryArray::try_from_iter(vals.map(|i| i.to_le_bytes())).unwrap();
687            let keys = Int8Array::from_iter_values(0..64);
688            let arr = DictionaryArray::try_new(keys, Arc::new(dict)).unwrap();
689            RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap()
690        };
691        let out = concat_batches(&schema, &[make(0..64), make(64..128)]).unwrap();
692        assert_eq!(out.num_rows(), 128);
693        let dict = out.column(0).as_dictionary::<Int8Type>();
694        assert_eq!(dict.values().len(), 128);
695    }
696
697    #[test]
698    fn test_concat_empty_vec() {
699        let re = concat(&[]);
700        assert!(re.is_err());
701    }
702
703    #[test]
704    fn test_concat_batches_no_columns() {
705        // Test concat using empty schema / batches without columns
706        let schema = Arc::new(Schema::empty());
707
708        let mut options = RecordBatchOptions::default();
709        options.row_count = Some(100);
710        let batch = RecordBatch::try_new_with_options(schema.clone(), vec![], &options).unwrap();
711        // put in 2 batches of 100 rows each
712        let re = concat_batches(&schema, &[batch.clone(), batch]).unwrap();
713
714        assert_eq!(re.num_rows(), 200);
715    }
716
717    #[test]
718    fn test_concat_one_element_vec() {
719        let arr = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
720            Some(-1),
721            Some(2),
722            None,
723        ])) as ArrayRef;
724        let result = concat(&[arr.as_ref()]).unwrap();
725        assert_eq!(
726            &arr, &result,
727            "concatenating single element array gives back the same result"
728        );
729    }
730
731    #[test]
732    fn test_concat_incompatible_datatypes() {
733        let re = concat(&[
734            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(2), None]),
735            // 2 string to make sure we only mention unique types
736            &StringArray::from(vec![Some("hello"), Some("bar"), Some("world")]),
737            &StringArray::from(vec![Some("hey"), Some(""), Some("you")]),
738            // Another type to make sure we are showing all the incompatible types
739            &PrimitiveArray::<Int32Type>::from(vec![Some(-1), Some(2), None]),
740        ]);
741
742        assert_eq!(
743            re.unwrap_err().to_string(),
744            "Invalid argument error: It is not possible to concatenate arrays of different data types (Int64, Utf8, Int32)."
745        );
746    }
747
748    #[test]
749    fn test_concat_10_incompatible_datatypes_should_include_all_of_them() {
750        let re = concat(&[
751            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(2), None]),
752            // 2 string to make sure we only mention unique types
753            &StringArray::from(vec![Some("hello"), Some("bar"), Some("world")]),
754            &StringArray::from(vec![Some("hey"), Some(""), Some("you")]),
755            // Another type to make sure we are showing all the incompatible types
756            &PrimitiveArray::<Int32Type>::from(vec![Some(-1), Some(2), None]),
757            &PrimitiveArray::<Int8Type>::from(vec![Some(-1), Some(2), None]),
758            &PrimitiveArray::<Int16Type>::from(vec![Some(-1), Some(2), None]),
759            &PrimitiveArray::<UInt8Type>::from(vec![Some(1), Some(2), None]),
760            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
761            &PrimitiveArray::<UInt32Type>::from(vec![Some(1), Some(2), None]),
762            // Non unique
763            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
764            &PrimitiveArray::<UInt64Type>::from(vec![Some(1), Some(2), None]),
765            &PrimitiveArray::<Float32Type>::from(vec![Some(1.0), Some(2.0), None]),
766        ]);
767
768        assert_eq!(
769            re.unwrap_err().to_string(),
770            "Invalid argument error: It is not possible to concatenate arrays of different data types (Int64, Utf8, Int32, Int8, Int16, UInt8, UInt16, UInt32, UInt64, Float32)."
771        );
772    }
773
774    #[test]
775    fn test_concat_11_incompatible_datatypes_should_only_include_10() {
776        let re = concat(&[
777            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(2), None]),
778            // 2 string to make sure we only mention unique types
779            &StringArray::from(vec![Some("hello"), Some("bar"), Some("world")]),
780            &StringArray::from(vec![Some("hey"), Some(""), Some("you")]),
781            // Another type to make sure we are showing all the incompatible types
782            &PrimitiveArray::<Int32Type>::from(vec![Some(-1), Some(2), None]),
783            &PrimitiveArray::<Int8Type>::from(vec![Some(-1), Some(2), None]),
784            &PrimitiveArray::<Int16Type>::from(vec![Some(-1), Some(2), None]),
785            &PrimitiveArray::<UInt8Type>::from(vec![Some(1), Some(2), None]),
786            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
787            &PrimitiveArray::<UInt32Type>::from(vec![Some(1), Some(2), None]),
788            // Non unique
789            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
790            &PrimitiveArray::<UInt64Type>::from(vec![Some(1), Some(2), None]),
791            &PrimitiveArray::<Float32Type>::from(vec![Some(1.0), Some(2.0), None]),
792            &PrimitiveArray::<Float64Type>::from(vec![Some(1.0), Some(2.0), None]),
793        ]);
794
795        assert_eq!(
796            re.unwrap_err().to_string(),
797            "Invalid argument error: It is not possible to concatenate arrays of different data types (Int64, Utf8, Int32, Int8, Int16, UInt8, UInt16, UInt32, UInt64, Float32, ...)."
798        );
799    }
800
801    #[test]
802    #[cfg_attr(miri, ignore)] // Unsupported inline assembly
803    fn test_concat_13_incompatible_datatypes_should_not_include_all_of_them() {
804        let re = concat(&[
805            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(2), None]),
806            // 2 string to make sure we only mention unique types
807            &StringArray::from(vec![Some("hello"), Some("bar"), Some("world")]),
808            &StringArray::from(vec![Some("hey"), Some(""), Some("you")]),
809            // Another type to make sure we are showing all the incompatible types
810            &PrimitiveArray::<Int32Type>::from(vec![Some(-1), Some(2), None]),
811            &PrimitiveArray::<Int8Type>::from(vec![Some(-1), Some(2), None]),
812            &PrimitiveArray::<Int16Type>::from(vec![Some(-1), Some(2), None]),
813            &PrimitiveArray::<UInt8Type>::from(vec![Some(1), Some(2), None]),
814            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
815            &PrimitiveArray::<UInt32Type>::from(vec![Some(1), Some(2), None]),
816            // Non unique
817            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
818            &PrimitiveArray::<UInt64Type>::from(vec![Some(1), Some(2), None]),
819            &PrimitiveArray::<Float32Type>::from(vec![Some(1.0), Some(2.0), None]),
820            &PrimitiveArray::<Float64Type>::from(vec![Some(1.0), Some(2.0), None]),
821            &PrimitiveArray::<Float16Type>::new_null(3),
822            &BooleanArray::from(vec![Some(true), Some(false), None]),
823        ]);
824
825        assert_eq!(
826            re.unwrap_err().to_string(),
827            "Invalid argument error: It is not possible to concatenate arrays of different data types (Int64, Utf8, Int32, Int8, Int16, UInt8, UInt16, UInt32, UInt64, Float32, ...)."
828        );
829    }
830
831    #[test]
832    fn test_concat_string_arrays() {
833        let arr = concat(&[
834            &StringArray::from(vec!["hello", "world"]),
835            &StringArray::from(vec!["2", "3", "4"]),
836            &StringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]),
837        ])
838        .unwrap();
839
840        let expected_output = Arc::new(StringArray::from(vec![
841            Some("hello"),
842            Some("world"),
843            Some("2"),
844            Some("3"),
845            Some("4"),
846            Some("foo"),
847            Some("bar"),
848            None,
849            Some("baz"),
850        ])) as ArrayRef;
851
852        assert_eq!(&arr, &expected_output);
853    }
854
855    #[test]
856    fn test_concat_string_view_arrays() {
857        let arr = concat(&[
858            &StringViewArray::from(vec!["helloxxxxxxxxxxa", "world____________"]),
859            &StringViewArray::from(vec!["helloxxxxxxxxxxy", "3", "4"]),
860            &StringViewArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]),
861        ])
862        .unwrap();
863
864        let expected_output = Arc::new(StringViewArray::from(vec![
865            Some("helloxxxxxxxxxxa"),
866            Some("world____________"),
867            Some("helloxxxxxxxxxxy"),
868            Some("3"),
869            Some("4"),
870            Some("foo"),
871            Some("bar"),
872            None,
873            Some("baz"),
874        ])) as ArrayRef;
875
876        assert_eq!(&arr, &expected_output);
877    }
878
879    #[test]
880    fn test_concat_primitive_arrays() {
881        let arr = concat(&[
882            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(-1), Some(2), None, None]),
883            &PrimitiveArray::<Int64Type>::from(vec![Some(101), Some(102), Some(103), None]),
884            &PrimitiveArray::<Int64Type>::from(vec![Some(256), Some(512), Some(1024)]),
885        ])
886        .unwrap();
887
888        let expected_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
889            Some(-1),
890            Some(-1),
891            Some(2),
892            None,
893            None,
894            Some(101),
895            Some(102),
896            Some(103),
897            None,
898            Some(256),
899            Some(512),
900            Some(1024),
901        ])) as ArrayRef;
902
903        assert_eq!(&arr, &expected_output);
904    }
905
906    #[test]
907    fn test_concat_primitive_array_slices() {
908        let input_1 =
909            PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(-1), Some(2), None, None])
910                .slice(1, 3);
911
912        let input_2 =
913            PrimitiveArray::<Int64Type>::from(vec![Some(101), Some(102), Some(103), None])
914                .slice(1, 3);
915        let arr = concat(&[&input_1, &input_2]).unwrap();
916
917        let expected_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
918            Some(-1),
919            Some(2),
920            None,
921            Some(102),
922            Some(103),
923            None,
924        ])) as ArrayRef;
925
926        assert_eq!(&arr, &expected_output);
927    }
928
929    #[test]
930    fn test_concat_boolean_primitive_arrays() {
931        let arr = concat(&[
932            &BooleanArray::from(vec![
933                Some(true),
934                Some(true),
935                Some(false),
936                None,
937                None,
938                Some(false),
939            ]),
940            &BooleanArray::from(vec![None, Some(false), Some(true), Some(false)]),
941        ])
942        .unwrap();
943
944        let expected_output = Arc::new(BooleanArray::from(vec![
945            Some(true),
946            Some(true),
947            Some(false),
948            None,
949            None,
950            Some(false),
951            None,
952            Some(false),
953            Some(true),
954            Some(false),
955        ])) as ArrayRef;
956
957        assert_eq!(&arr, &expected_output);
958    }
959
960    #[test]
961    fn test_concat_primitive_list_arrays() {
962        let list1 = [
963            Some(vec![Some(-1), Some(-1), Some(2), None, None]),
964            Some(vec![]),
965            None,
966            Some(vec![Some(10)]),
967        ];
968        let list1_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone());
969
970        let list2 = [
971            None,
972            Some(vec![Some(100), None, Some(101)]),
973            Some(vec![Some(102)]),
974        ];
975        let list2_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone());
976
977        let list3 = [Some(vec![Some(1000), Some(1001)])];
978        let list3_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list3.clone());
979
980        let array_result = concat(&[&list1_array, &list2_array, &list3_array]).unwrap();
981
982        let expected = list1.into_iter().chain(list2).chain(list3);
983        let array_expected = ListArray::from_iter_primitive::<Int64Type, _, _>(expected);
984
985        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
986    }
987
988    #[test]
989    fn test_concat_primitive_list_arrays_slices() {
990        let list1 = [
991            Some(vec![Some(-1), Some(-1), Some(2), None, None]),
992            Some(vec![]), // In slice
993            None,         // In slice
994            Some(vec![Some(10)]),
995        ];
996        let list1_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone());
997        let list1_array = list1_array.slice(1, 2);
998        let list1_values = list1.into_iter().skip(1).take(2);
999
1000        let list2 = [
1001            None,
1002            Some(vec![Some(100), None, Some(101)]),
1003            Some(vec![Some(102)]),
1004        ];
1005        let list2_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone());
1006
1007        // verify that this test covers the case when the first offset is non zero
1008        assert!(list1_array.offsets()[0].as_usize() > 0);
1009        let array_result = concat(&[&list1_array, &list2_array]).unwrap();
1010
1011        let expected = list1_values.chain(list2);
1012        let array_expected = ListArray::from_iter_primitive::<Int64Type, _, _>(expected);
1013
1014        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1015    }
1016
1017    #[test]
1018    fn test_concat_primitive_list_arrays_sliced_lengths() {
1019        let list1 = [
1020            Some(vec![Some(-1), Some(-1), Some(2), None, None]), // In slice
1021            Some(vec![]),                                        // In slice
1022            None,                                                // In slice
1023            Some(vec![Some(10)]),
1024        ];
1025        let list1_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone());
1026        let list1_array = list1_array.slice(0, 3); // no offset, but not all values
1027        let list1_values = list1.into_iter().take(3);
1028
1029        let list2 = [
1030            None,
1031            Some(vec![Some(100), None, Some(101)]),
1032            Some(vec![Some(102)]),
1033        ];
1034        let list2_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone());
1035
1036        // verify that this test covers the case when the first offset is zero, but the
1037        // last offset doesn't cover the entire array
1038        assert_eq!(list1_array.offsets()[0].as_usize(), 0);
1039        assert!(list1_array.offsets().last().unwrap().as_usize() < list1_array.values().len());
1040        let array_result = concat(&[&list1_array, &list2_array]).unwrap();
1041
1042        let expected = list1_values.chain(list2);
1043        let array_expected = ListArray::from_iter_primitive::<Int64Type, _, _>(expected);
1044
1045        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1046    }
1047
1048    #[test]
1049    fn test_concat_primitive_fixed_size_list_arrays() {
1050        let list1 = [
1051            Some(vec![Some(-1), None]),
1052            None,
1053            Some(vec![Some(10), Some(20)]),
1054        ];
1055        let list1_array =
1056            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone(), 2);
1057
1058        let list2 = [
1059            None,
1060            Some(vec![Some(100), None]),
1061            Some(vec![Some(102), Some(103)]),
1062        ];
1063        let list2_array =
1064            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone(), 2);
1065
1066        let list3 = [Some(vec![Some(1000), Some(1001)])];
1067        let list3_array =
1068            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(list3.clone(), 2);
1069
1070        let array_result = concat(&[&list1_array, &list2_array, &list3_array]).unwrap();
1071
1072        let expected = list1.into_iter().chain(list2).chain(list3);
1073        let array_expected =
1074            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(expected, 2);
1075
1076        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1077    }
1078
1079    #[test]
1080    fn test_concat_list_view_arrays() {
1081        let list1 = [
1082            Some(vec![Some(-1), None]),
1083            None,
1084            Some(vec![Some(10), Some(20)]),
1085        ];
1086        let mut list1_array = ListViewBuilder::new(Int64Builder::new());
1087        for v in list1.iter() {
1088            list1_array.append_option(v.clone());
1089        }
1090        let list1_array = list1_array.finish();
1091
1092        let list2 = [
1093            None,
1094            Some(vec![Some(100), None]),
1095            Some(vec![Some(102), Some(103)]),
1096        ];
1097        let mut list2_array = ListViewBuilder::new(Int64Builder::new());
1098        for v in list2.iter() {
1099            list2_array.append_option(v.clone());
1100        }
1101        let list2_array = list2_array.finish();
1102
1103        let list3 = [Some(vec![Some(1000), Some(1001)])];
1104        let mut list3_array = ListViewBuilder::new(Int64Builder::new());
1105        for v in list3.iter() {
1106            list3_array.append_option(v.clone());
1107        }
1108        let list3_array = list3_array.finish();
1109
1110        let array_result = concat(&[&list1_array, &list2_array, &list3_array]).unwrap();
1111
1112        let expected: Vec<_> = list1.into_iter().chain(list2).chain(list3).collect();
1113        let mut array_expected = ListViewBuilder::new(Int64Builder::new());
1114        for v in expected.iter() {
1115            array_expected.append_option(v.clone());
1116        }
1117        let array_expected = array_expected.finish();
1118
1119        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1120    }
1121
1122    #[test]
1123    fn test_concat_sliced_list_view_arrays() {
1124        let list1 = [
1125            Some(vec![Some(-1), None]),
1126            None,
1127            Some(vec![Some(10), Some(20)]),
1128        ];
1129        let mut list1_array = ListViewBuilder::new(Int64Builder::new());
1130        for v in list1.iter() {
1131            list1_array.append_option(v.clone());
1132        }
1133        let list1_array = list1_array.finish();
1134
1135        let list2 = [
1136            None,
1137            Some(vec![Some(100), None]),
1138            Some(vec![Some(102), Some(103)]),
1139        ];
1140        let mut list2_array = ListViewBuilder::new(Int64Builder::new());
1141        for v in list2.iter() {
1142            list2_array.append_option(v.clone());
1143        }
1144        let list2_array = list2_array.finish();
1145
1146        let list3 = [Some(vec![Some(1000), Some(1001)])];
1147        let mut list3_array = ListViewBuilder::new(Int64Builder::new());
1148        for v in list3.iter() {
1149            list3_array.append_option(v.clone());
1150        }
1151        let list3_array = list3_array.finish();
1152
1153        // Concat sliced arrays.
1154        // ListView slicing will slice the offset/sizes but preserve the original values child.
1155        let array_result = concat(&[
1156            &list1_array.slice(1, 2),
1157            &list2_array.slice(1, 2),
1158            &list3_array.slice(0, 1),
1159        ])
1160        .unwrap();
1161
1162        let expected: Vec<_> = vec![
1163            None,
1164            Some(vec![Some(10), Some(20)]),
1165            Some(vec![Some(100), None]),
1166            Some(vec![Some(102), Some(103)]),
1167            Some(vec![Some(1000), Some(1001)]),
1168        ];
1169        let mut array_expected = ListViewBuilder::new(Int64Builder::new());
1170        for v in expected.iter() {
1171            array_expected.append_option(v.clone());
1172        }
1173        let array_expected = array_expected.finish();
1174
1175        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1176    }
1177
1178    #[test]
1179    fn test_concat_struct_arrays() {
1180        let field = Arc::new(Field::new("field", DataType::Int64, true));
1181        let input_primitive_1: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1182            Some(-1),
1183            Some(-1),
1184            Some(2),
1185            None,
1186            None,
1187        ]));
1188        let input_struct_1 = StructArray::from(vec![(field.clone(), input_primitive_1)]);
1189
1190        let input_primitive_2: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1191            Some(101),
1192            Some(102),
1193            Some(103),
1194            None,
1195        ]));
1196        let input_struct_2 = StructArray::from(vec![(field.clone(), input_primitive_2)]);
1197
1198        let input_primitive_3: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1199            Some(256),
1200            Some(512),
1201            Some(1024),
1202        ]));
1203        let input_struct_3 = StructArray::from(vec![(field, input_primitive_3)]);
1204
1205        let arr = concat(&[&input_struct_1, &input_struct_2, &input_struct_3]).unwrap();
1206
1207        let expected_primitive_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1208            Some(-1),
1209            Some(-1),
1210            Some(2),
1211            None,
1212            None,
1213            Some(101),
1214            Some(102),
1215            Some(103),
1216            None,
1217            Some(256),
1218            Some(512),
1219            Some(1024),
1220        ])) as ArrayRef;
1221
1222        let actual_primitive = arr
1223            .as_any()
1224            .downcast_ref::<StructArray>()
1225            .unwrap()
1226            .column(0);
1227        assert_eq!(actual_primitive, &expected_primitive_output);
1228    }
1229
1230    #[test]
1231    fn test_concat_struct_array_slices() {
1232        let field = Arc::new(Field::new("field", DataType::Int64, true));
1233        let input_primitive_1: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1234            Some(-1),
1235            Some(-1),
1236            Some(2),
1237            None,
1238            None,
1239        ]));
1240        let input_struct_1 = StructArray::from(vec![(field.clone(), input_primitive_1)]);
1241
1242        let input_primitive_2: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1243            Some(101),
1244            Some(102),
1245            Some(103),
1246            None,
1247        ]));
1248        let input_struct_2 = StructArray::from(vec![(field, input_primitive_2)]);
1249
1250        let arr = concat(&[&input_struct_1.slice(1, 3), &input_struct_2.slice(1, 2)]).unwrap();
1251
1252        let expected_primitive_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1253            Some(-1),
1254            Some(2),
1255            None,
1256            Some(102),
1257            Some(103),
1258        ])) as ArrayRef;
1259
1260        let actual_primitive = arr
1261            .as_any()
1262            .downcast_ref::<StructArray>()
1263            .unwrap()
1264            .column(0);
1265        assert_eq!(actual_primitive, &expected_primitive_output);
1266    }
1267
1268    #[test]
1269    fn test_concat_struct_arrays_no_nulls() {
1270        let input_1a = vec![1, 2, 3];
1271        let input_1b = vec!["one", "two", "three"];
1272        let input_2a = vec![4, 5, 6, 7];
1273        let input_2b = vec!["four", "five", "six", "seven"];
1274
1275        let struct_from_primitives = |ints: Vec<i64>, strings: Vec<&str>| {
1276            StructArray::try_from(vec![
1277                ("ints", Arc::new(Int64Array::from(ints)) as _),
1278                ("strings", Arc::new(StringArray::from(strings)) as _),
1279            ])
1280        };
1281
1282        let expected_output = struct_from_primitives(
1283            [input_1a.clone(), input_2a.clone()].concat(),
1284            [input_1b.clone(), input_2b.clone()].concat(),
1285        )
1286        .unwrap();
1287
1288        let input_1 = struct_from_primitives(input_1a, input_1b).unwrap();
1289        let input_2 = struct_from_primitives(input_2a, input_2b).unwrap();
1290
1291        let arr = concat(&[&input_1, &input_2]).unwrap();
1292        let struct_result = arr.as_struct();
1293
1294        assert_eq!(struct_result, &expected_output);
1295        assert_eq!(arr.null_count(), 0);
1296    }
1297
1298    #[test]
1299    fn test_concat_struct_no_fields() {
1300        let input_1 = StructArray::new_empty_fields(10, None);
1301        let input_2 = StructArray::new_empty_fields(10, None);
1302        let arr = concat(&[&input_1, &input_2]).unwrap();
1303
1304        assert_eq!(arr.len(), 20);
1305        assert_eq!(arr.null_count(), 0);
1306
1307        let input1_valid = StructArray::new_empty_fields(10, Some(NullBuffer::new_valid(10)));
1308        let input2_null = StructArray::new_empty_fields(10, Some(NullBuffer::new_null(10)));
1309        let arr = concat(&[&input1_valid, &input2_null]).unwrap();
1310
1311        assert_eq!(arr.len(), 20);
1312        assert_eq!(arr.null_count(), 10);
1313    }
1314
1315    #[test]
1316    fn test_string_array_slices() {
1317        let input_1 = StringArray::from(vec!["hello", "A", "B", "C"]);
1318        let input_2 = StringArray::from(vec!["world", "D", "E", "Z"]);
1319
1320        let arr = concat(&[&input_1.slice(1, 3), &input_2.slice(1, 2)]).unwrap();
1321
1322        let expected_output = StringArray::from(vec!["A", "B", "C", "D", "E"]);
1323
1324        let actual_output = arr.as_any().downcast_ref::<StringArray>().unwrap();
1325        assert_eq!(actual_output, &expected_output);
1326    }
1327
1328    #[test]
1329    fn test_string_array_with_null_slices() {
1330        let input_1 = StringArray::from(vec![Some("hello"), None, Some("A"), Some("C")]);
1331        let input_2 = StringArray::from(vec![None, Some("world"), Some("D"), None]);
1332
1333        let arr = concat(&[&input_1.slice(1, 3), &input_2.slice(1, 2)]).unwrap();
1334
1335        let expected_output =
1336            StringArray::from(vec![None, Some("A"), Some("C"), Some("world"), Some("D")]);
1337
1338        let actual_output = arr.as_any().downcast_ref::<StringArray>().unwrap();
1339        assert_eq!(actual_output, &expected_output);
1340    }
1341
1342    fn collect_string_dictionary(array: &DictionaryArray<Int32Type>) -> Vec<Option<&str>> {
1343        let concrete = array.downcast_dict::<StringArray>().unwrap();
1344        concrete.into_iter().collect()
1345    }
1346
1347    #[test]
1348    fn test_string_dictionary_array() {
1349        let input_1: DictionaryArray<Int32Type> = vec!["hello", "A", "B", "hello", "hello", "C"]
1350            .into_iter()
1351            .collect();
1352        let input_2: DictionaryArray<Int32Type> = vec!["hello", "E", "E", "hello", "F", "E"]
1353            .into_iter()
1354            .collect();
1355
1356        let expected: Vec<_> = vec![
1357            "hello", "A", "B", "hello", "hello", "C", "hello", "E", "E", "hello", "F", "E",
1358        ]
1359        .into_iter()
1360        .map(Some)
1361        .collect();
1362
1363        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1364        let dictionary = concat.as_dictionary::<Int32Type>();
1365        let actual = collect_string_dictionary(dictionary);
1366        assert_eq!(actual, expected);
1367
1368        // Should have concatenated inputs together
1369        assert_eq!(
1370            dictionary.values().len(),
1371            input_1.values().len() + input_2.values().len(),
1372        )
1373    }
1374
1375    #[test]
1376    fn test_string_dictionary_array_nulls() {
1377        let input_1: DictionaryArray<Int32Type> = vec![Some("foo"), Some("bar"), None, Some("fiz")]
1378            .into_iter()
1379            .collect();
1380        let input_2: DictionaryArray<Int32Type> = vec![None].into_iter().collect();
1381        let expected = vec![Some("foo"), Some("bar"), None, Some("fiz"), None];
1382
1383        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1384        let dictionary = concat.as_dictionary::<Int32Type>();
1385        let actual = collect_string_dictionary(dictionary);
1386        assert_eq!(actual, expected);
1387
1388        // Should have concatenated inputs together
1389        assert_eq!(
1390            dictionary.values().len(),
1391            input_1.values().len() + input_2.values().len(),
1392        )
1393    }
1394
1395    #[test]
1396    fn test_string_dictionary_array_nulls_in_values() {
1397        let input_1_keys = Int32Array::from_iter_values([0, 2, 1, 3]);
1398        let input_1_values = StringArray::from(vec![Some("foo"), None, Some("bar"), Some("fiz")]);
1399        let input_1 = DictionaryArray::new(input_1_keys, Arc::new(input_1_values));
1400
1401        let input_2_keys = Int32Array::from_iter_values([0]);
1402        let input_2_values = StringArray::from(vec![None, Some("hello")]);
1403        let input_2 = DictionaryArray::new(input_2_keys, Arc::new(input_2_values));
1404
1405        let expected = vec![Some("foo"), Some("bar"), None, Some("fiz"), None];
1406
1407        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1408        let dictionary = concat.as_dictionary::<Int32Type>();
1409        let actual = collect_string_dictionary(dictionary);
1410        assert_eq!(actual, expected);
1411    }
1412
1413    #[test]
1414    fn test_string_dictionary_merge() {
1415        let mut builder = StringDictionaryBuilder::<Int32Type>::new();
1416        for i in 0..20 {
1417            builder.append(i.to_string()).unwrap();
1418        }
1419        let input_1 = builder.finish();
1420
1421        let mut builder = StringDictionaryBuilder::<Int32Type>::new();
1422        for i in 0..30 {
1423            builder.append(i.to_string()).unwrap();
1424        }
1425        let input_2 = builder.finish();
1426
1427        let expected: Vec<_> = (0..20).chain(0..30).map(|x| x.to_string()).collect();
1428        let expected: Vec<_> = expected.iter().map(|x| Some(x.as_str())).collect();
1429
1430        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1431        let dictionary = concat.as_dictionary::<Int32Type>();
1432        let actual = collect_string_dictionary(dictionary);
1433        assert_eq!(actual, expected);
1434
1435        // Should have merged inputs together
1436        // Not 30 as this is done on a best-effort basis
1437        let values_len = dictionary.values().len();
1438        assert!((30..40).contains(&values_len), "{values_len}")
1439    }
1440
1441    #[test]
1442    fn test_primitive_dictionary_merge() {
1443        // Same value repeated 5 times.
1444        let keys = vec![1; 5];
1445        let values = (10..20).collect::<Vec<_>>();
1446        let dict = DictionaryArray::new(
1447            Int8Array::from(keys.clone()),
1448            Arc::new(Int32Array::from(values.clone())),
1449        );
1450        let other = DictionaryArray::new(
1451            Int8Array::from(keys.clone()),
1452            Arc::new(Int32Array::from(values.clone())),
1453        );
1454
1455        let result_same_dictionary = concat(&[&dict, &dict]).unwrap();
1456        // Verify pointer equality check succeeds, and therefore the
1457        // dictionaries are not merged. A single values buffer should be reused
1458        // in this case.
1459        assert!(
1460            dict.values().to_data().ptr_eq(
1461                &result_same_dictionary
1462                    .as_dictionary::<Int8Type>()
1463                    .values()
1464                    .to_data()
1465            )
1466        );
1467        assert_eq!(
1468            result_same_dictionary
1469                .as_dictionary::<Int8Type>()
1470                .values()
1471                .len(),
1472            values.len(),
1473        );
1474
1475        let result_cloned_dictionary = concat(&[&dict, &other]).unwrap();
1476        // Should have only 1 underlying value since all keys reference it.
1477        assert_eq!(
1478            result_cloned_dictionary
1479                .as_dictionary::<Int8Type>()
1480                .values()
1481                .len(),
1482            1
1483        );
1484    }
1485
1486    #[test]
1487    fn test_concat_string_sizes() {
1488        let a: LargeStringArray = ((0..150).map(|_| Some("foo"))).collect();
1489        let b: LargeStringArray = ((0..150).map(|_| Some("foo"))).collect();
1490        let c = LargeStringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]);
1491        // 150 * 3 = 450
1492        // 150 * 3 = 450
1493        // 3 * 3   = 9
1494        // ------------+
1495        // 909
1496
1497        let arr = concat(&[&a, &b, &c]).unwrap();
1498        assert_eq!(arr.to_data().buffers()[1].capacity(), 909);
1499    }
1500
1501    #[test]
1502    fn test_dictionary_concat_reuse() {
1503        let array: DictionaryArray<Int8Type> = vec!["a", "a", "b", "c"].into_iter().collect();
1504        let copy: DictionaryArray<Int8Type> = array.clone();
1505
1506        // dictionary is "a", "b", "c"
1507        assert_eq!(
1508            array.values(),
1509            &(Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef)
1510        );
1511        assert_eq!(array.keys(), &Int8Array::from(vec![0, 0, 1, 2]));
1512
1513        // concatenate it with itself
1514        let combined = concat(&[&copy as _, &array as _]).unwrap();
1515        let combined = combined.as_dictionary::<Int8Type>();
1516
1517        assert_eq!(
1518            combined.values(),
1519            &(Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef),
1520            "Actual: {combined:#?}"
1521        );
1522
1523        assert_eq!(
1524            combined.keys(),
1525            &Int8Array::from(vec![0, 0, 1, 2, 0, 0, 1, 2])
1526        );
1527
1528        // Should have reused the dictionary
1529        assert!(
1530            array
1531                .values()
1532                .to_data()
1533                .ptr_eq(&combined.values().to_data())
1534        );
1535        assert!(copy.values().to_data().ptr_eq(&combined.values().to_data()));
1536
1537        let new: DictionaryArray<Int8Type> = vec!["d"].into_iter().collect();
1538        let combined = concat(&[&copy as _, &array as _, &new as _]).unwrap();
1539        let com = combined.as_dictionary::<Int8Type>();
1540
1541        // Should not have reused the dictionary
1542        assert!(!array.values().to_data().ptr_eq(&com.values().to_data()));
1543        assert!(!copy.values().to_data().ptr_eq(&com.values().to_data()));
1544        assert!(!new.values().to_data().ptr_eq(&com.values().to_data()));
1545    }
1546
1547    #[test]
1548    fn concat_record_batches() {
1549        let schema = Arc::new(Schema::new(vec![
1550            Field::new("a", DataType::Int32, false),
1551            Field::new("b", DataType::Utf8, false),
1552        ]));
1553        let batch1 = RecordBatch::try_new(
1554            schema.clone(),
1555            vec![
1556                Arc::new(Int32Array::from(vec![1, 2])),
1557                Arc::new(StringArray::from(vec!["a", "b"])),
1558            ],
1559        )
1560        .unwrap();
1561        let batch2 = RecordBatch::try_new(
1562            schema.clone(),
1563            vec![
1564                Arc::new(Int32Array::from(vec![3, 4])),
1565                Arc::new(StringArray::from(vec!["c", "d"])),
1566            ],
1567        )
1568        .unwrap();
1569        let new_batch = concat_batches(&schema, [&batch1, &batch2]).unwrap();
1570        assert_eq!(new_batch.schema().as_ref(), schema.as_ref());
1571        assert_eq!(2, new_batch.num_columns());
1572        assert_eq!(4, new_batch.num_rows());
1573        let new_batch_owned = concat_batches(&schema, &[batch1, batch2]).unwrap();
1574        assert_eq!(new_batch_owned.schema().as_ref(), schema.as_ref());
1575        assert_eq!(2, new_batch_owned.num_columns());
1576        assert_eq!(4, new_batch_owned.num_rows());
1577    }
1578
1579    #[test]
1580    fn concat_empty_record_batch() {
1581        let schema = Arc::new(Schema::new(vec![
1582            Field::new("a", DataType::Int32, false),
1583            Field::new("b", DataType::Utf8, false),
1584        ]));
1585        let batch = concat_batches(&schema, []).unwrap();
1586        assert_eq!(batch.schema().as_ref(), schema.as_ref());
1587        assert_eq!(0, batch.num_rows());
1588    }
1589
1590    #[test]
1591    fn concat_record_batches_of_different_schemas_but_compatible_data() {
1592        let schema1 = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
1593        // column names differ
1594        let schema2 = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, false)]));
1595        let batch1 = RecordBatch::try_new(
1596            schema1.clone(),
1597            vec![Arc::new(Int32Array::from(vec![1, 2]))],
1598        )
1599        .unwrap();
1600        let batch2 =
1601            RecordBatch::try_new(schema2, vec![Arc::new(Int32Array::from(vec![3, 4]))]).unwrap();
1602        // concat_batches simply uses the schema provided
1603        let batch = concat_batches(&schema1, [&batch1, &batch2]).unwrap();
1604        assert_eq!(batch.schema().as_ref(), schema1.as_ref());
1605        assert_eq!(4, batch.num_rows());
1606    }
1607
1608    #[test]
1609    fn concat_record_batches_of_different_schemas_incompatible_data() {
1610        let schema1 = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
1611        // column names differ
1612        let schema2 = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, false)]));
1613        let batch1 = RecordBatch::try_new(
1614            schema1.clone(),
1615            vec![Arc::new(Int32Array::from(vec![1, 2]))],
1616        )
1617        .unwrap();
1618        let batch2 = RecordBatch::try_new(
1619            schema2,
1620            vec![Arc::new(StringArray::from(vec!["foo", "bar"]))],
1621        )
1622        .unwrap();
1623
1624        let error = concat_batches(&schema1, [&batch1, &batch2]).unwrap_err();
1625        assert_eq!(
1626            error.to_string(),
1627            "Invalid argument error: It is not possible to concatenate arrays of different data types (Int32, Utf8)."
1628        );
1629    }
1630
1631    #[test]
1632    fn concat_capacity() {
1633        let a = Int32Array::from_iter_values(0..100);
1634        let b = Int32Array::from_iter_values(10..20);
1635        let a = concat(&[&a, &b]).unwrap();
1636        let data = a.to_data();
1637        assert_eq!(data.buffers()[0].len(), 440);
1638        assert_eq!(data.buffers()[0].capacity(), 440);
1639
1640        let a = concat(&[&a.slice(10, 20), &b]).unwrap();
1641        let data = a.to_data();
1642        assert_eq!(data.buffers()[0].len(), 120);
1643        assert_eq!(data.buffers()[0].capacity(), 120);
1644
1645        let a = StringArray::from_iter_values(std::iter::repeat_n("foo", 100));
1646        let b = StringArray::from(vec!["bingo", "bongo", "lorem", ""]);
1647
1648        let a = concat(&[&a, &b]).unwrap();
1649        let data = a.to_data();
1650        // (100 + 4 + 1) * size_of<i32>()
1651        assert_eq!(data.buffers()[0].len(), 420);
1652        assert_eq!(data.buffers()[0].capacity(), 420);
1653
1654        // len("foo") * 100 + len("bingo") + len("bongo") + len("lorem")
1655        assert_eq!(data.buffers()[1].len(), 315);
1656        assert_eq!(data.buffers()[1].capacity(), 315);
1657
1658        let a = concat(&[&a.slice(10, 40), &b]).unwrap();
1659        let data = a.to_data();
1660        // (40 + 4 + 5) * size_of<i32>()
1661        assert_eq!(data.buffers()[0].len(), 180);
1662        assert_eq!(data.buffers()[0].capacity(), 180);
1663
1664        // len("foo") * 40 + len("bingo") + len("bongo") + len("lorem")
1665        assert_eq!(data.buffers()[1].len(), 135);
1666        assert_eq!(data.buffers()[1].capacity(), 135);
1667
1668        let a = LargeBinaryArray::from_iter_values(std::iter::repeat_n(b"foo", 100));
1669        let b = LargeBinaryArray::from_iter_values(std::iter::repeat_n(b"cupcakes", 10));
1670
1671        let a = concat(&[&a, &b]).unwrap();
1672        let data = a.to_data();
1673        // (100 + 10 + 1) * size_of<i64>()
1674        assert_eq!(data.buffers()[0].len(), 888);
1675        assert_eq!(data.buffers()[0].capacity(), 888);
1676
1677        // len("foo") * 100 + len("cupcakes") * 10
1678        assert_eq!(data.buffers()[1].len(), 380);
1679        assert_eq!(data.buffers()[1].capacity(), 380);
1680
1681        let a = concat(&[&a.slice(10, 40), &b]).unwrap();
1682        let data = a.to_data();
1683        // (40 + 10 + 1) * size_of<i64>()
1684        assert_eq!(data.buffers()[0].len(), 408);
1685        assert_eq!(data.buffers()[0].capacity(), 408);
1686
1687        // len("foo") * 40 + len("cupcakes") * 10
1688        assert_eq!(data.buffers()[1].len(), 200);
1689        assert_eq!(data.buffers()[1].capacity(), 200);
1690    }
1691
1692    #[test]
1693    fn concat_sparse_nulls() {
1694        let values = StringArray::from_iter_values((0..100).map(|x| x.to_string()));
1695        let keys = Int32Array::from(vec![1; 10]);
1696        let dict_a = DictionaryArray::new(keys, Arc::new(values));
1697        let values = StringArray::new_null(0);
1698        let keys = Int32Array::new_null(10);
1699        let dict_b = DictionaryArray::new(keys, Arc::new(values));
1700        let array = concat(&[&dict_a, &dict_b]).unwrap();
1701        assert_eq!(array.null_count(), 10);
1702        assert_eq!(array.logical_null_count(), 10);
1703    }
1704
1705    #[test]
1706    fn concat_dictionary_list_array_simple() {
1707        let scalars = [
1708            create_single_row_list_of_dict(vec![Some("a")]),
1709            create_single_row_list_of_dict(vec![Some("a")]),
1710            create_single_row_list_of_dict(vec![Some("b")]),
1711        ];
1712
1713        let arrays = scalars.iter().map(|a| a as &dyn Array).collect::<Vec<_>>();
1714        let concat_res = concat(arrays.as_slice()).unwrap();
1715
1716        let expected_list = create_list_of_dict(vec![
1717            // Row 1
1718            Some(vec![Some("a")]),
1719            Some(vec![Some("a")]),
1720            Some(vec![Some("b")]),
1721        ]);
1722
1723        let list = concat_res.as_list::<i32>();
1724
1725        // Assert that the list is equal to the expected list
1726        list.iter().zip(expected_list.iter()).for_each(|(a, b)| {
1727            assert_eq!(a, b);
1728        });
1729
1730        assert_dictionary_has_unique_values::<_, StringArray>(
1731            list.values().as_dictionary::<Int32Type>(),
1732        );
1733    }
1734
1735    #[test]
1736    #[cfg_attr(miri, ignore)] // Takes too long
1737    fn concat_many_dictionary_list_arrays() {
1738        let number_of_unique_values = 8;
1739        let scalars = (0..80000)
1740            .map(|i| {
1741                create_single_row_list_of_dict(vec![Some(
1742                    (i % number_of_unique_values).to_string(),
1743                )])
1744            })
1745            .collect::<Vec<_>>();
1746
1747        let arrays = scalars.iter().map(|a| a as &dyn Array).collect::<Vec<_>>();
1748        let concat_res = concat(arrays.as_slice()).unwrap();
1749
1750        let expected_list = create_list_of_dict(
1751            (0..80000)
1752                .map(|i| Some(vec![Some((i % number_of_unique_values).to_string())]))
1753                .collect::<Vec<_>>(),
1754        );
1755
1756        let list = concat_res.as_list::<i32>();
1757
1758        // Assert that the list is equal to the expected list
1759        list.iter().zip(expected_list.iter()).for_each(|(a, b)| {
1760            assert_eq!(a, b);
1761        });
1762
1763        assert_dictionary_has_unique_values::<_, StringArray>(
1764            list.values().as_dictionary::<Int32Type>(),
1765        );
1766    }
1767
1768    fn create_single_row_list_of_dict(
1769        list_items: Vec<Option<impl AsRef<str>>>,
1770    ) -> GenericListArray<i32> {
1771        let rows = list_items.into_iter().map(Some).collect();
1772
1773        create_list_of_dict(vec![rows])
1774    }
1775
1776    fn create_list_of_dict(
1777        rows: Vec<Option<Vec<Option<impl AsRef<str>>>>>,
1778    ) -> GenericListArray<i32> {
1779        let mut builder =
1780            GenericListBuilder::<i32, _>::new(StringDictionaryBuilder::<Int32Type>::new());
1781
1782        for row in rows {
1783            builder.append_option(row);
1784        }
1785
1786        builder.finish()
1787    }
1788
1789    fn assert_dictionary_has_unique_values<'a, K, V>(array: &'a DictionaryArray<K>)
1790    where
1791        K: ArrowDictionaryKeyType,
1792        V: Sync + Send + 'static,
1793        &'a V: ArrayAccessor + IntoIterator,
1794        <&'a V as ArrayAccessor>::Item: Default + Clone + PartialEq + Debug + Ord,
1795        <&'a V as IntoIterator>::Item: Clone + PartialEq + Debug + Ord,
1796    {
1797        let dict = array.downcast_dict::<V>().unwrap();
1798        let mut values = dict.values().into_iter().collect::<Vec<_>>();
1799
1800        // remove duplicates must be sorted first so we can compare
1801        values.sort();
1802
1803        let mut unique_values = values.clone();
1804
1805        unique_values.dedup();
1806
1807        assert_eq!(
1808            values, unique_values,
1809            "There are duplicates in the value list (the value list here is sorted which is only for the assertion)"
1810        );
1811    }
1812
1813    // Test the simple case of concatenating two RunArrays
1814    #[test]
1815    fn test_concat_run_array() {
1816        // Create simple run arrays
1817        let run_ends1 = Int32Array::from(vec![2, 4]);
1818        let values1 = Int32Array::from(vec![10, 20]);
1819        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1820
1821        let run_ends2 = Int32Array::from(vec![1, 4]);
1822        let values2 = Int32Array::from(vec![30, 40]);
1823        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
1824
1825        // Concatenate the arrays - this should now work properly
1826        let result = concat(&[&array1, &array2]).unwrap();
1827        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1828
1829        // Check that the result has the correct length
1830        assert_eq!(result_run_array.len(), 8); // 4 + 4
1831
1832        // Check the run ends
1833        let run_ends = result_run_array.run_ends().values();
1834        assert_eq!(run_ends.len(), 4);
1835        assert_eq!(&[2, 4, 5, 8], run_ends);
1836
1837        // Check the values
1838        let values = result_run_array
1839            .values()
1840            .as_any()
1841            .downcast_ref::<Int32Array>()
1842            .unwrap();
1843        assert_eq!(values.len(), 4);
1844        assert_eq!(&[10, 20, 30, 40], values.values());
1845    }
1846
1847    #[test]
1848    fn test_concat_sliced_run_array() {
1849        // Slicing away first run in both arrays
1850        let run_ends1 = Int32Array::from(vec![2, 4]);
1851        let values1 = Int32Array::from(vec![10, 20]);
1852        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap(); // [10, 10, 20, 20]
1853        let array1 = array1.slice(2, 2); // [20, 20]
1854
1855        let run_ends2 = Int32Array::from(vec![1, 4]);
1856        let values2 = Int32Array::from(vec![30, 40]);
1857        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap(); // [30, 40, 40, 40]
1858        let array2 = array2.slice(1, 3); // [40, 40, 40]
1859
1860        let result = concat(&[&array1, &array2]).unwrap();
1861        let result = result.as_run::<Int32Type>();
1862        let result = result.downcast::<Int32Array>().unwrap();
1863
1864        let expected = vec![20, 20, 40, 40, 40];
1865        let actual = result.into_iter().flatten().collect::<Vec<_>>();
1866        assert_eq!(expected, actual);
1867    }
1868
1869    #[test]
1870    fn test_concat_run_array_matching_first_last_value() {
1871        // Create a run array with run ends [2, 4, 7] and values [10, 20, 30]
1872        let run_ends1 = Int32Array::from(vec![2, 4, 7]);
1873        let values1 = Int32Array::from(vec![10, 20, 30]);
1874        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1875
1876        // Create another run array with run ends [3, 5] and values [30, 40]
1877        let run_ends2 = Int32Array::from(vec![3, 5]);
1878        let values2 = Int32Array::from(vec![30, 40]);
1879        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
1880
1881        // Concatenate the two arrays
1882        let result = concat(&[&array1, &array2]).unwrap();
1883        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1884
1885        // The result should have length 12 (7 + 5)
1886        assert_eq!(result_run_array.len(), 12);
1887
1888        // Check that the run ends are correct
1889        let run_ends = result_run_array.run_ends().values();
1890        assert_eq!(&[2, 4, 7, 10, 12], run_ends);
1891
1892        // Check that the values are correct
1893        assert_eq!(
1894            &[10, 20, 30, 30, 40],
1895            result_run_array
1896                .values()
1897                .as_any()
1898                .downcast_ref::<Int32Array>()
1899                .unwrap()
1900                .values()
1901        );
1902    }
1903
1904    #[test]
1905    fn test_concat_run_array_with_nulls() {
1906        // Create values array with nulls
1907        let values1 = Int32Array::from(vec![Some(10), None, Some(30)]);
1908        let run_ends1 = Int32Array::from(vec![2, 4, 7]);
1909        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1910
1911        // Create another run array with run ends [3, 5] and values [30, null]
1912        let values2 = Int32Array::from(vec![Some(30), None]);
1913        let run_ends2 = Int32Array::from(vec![3, 5]);
1914        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
1915
1916        // Concatenate the two arrays
1917        let result = concat(&[&array1, &array2]).unwrap();
1918        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1919
1920        // The result should have length 12 (7 + 5)
1921        assert_eq!(result_run_array.len(), 12);
1922
1923        // Get a reference to the run array itself for testing
1924
1925        // Just test the length and run ends without asserting specific values
1926        // This ensures the test passes while we work on full support for RunArray nulls
1927        assert_eq!(result_run_array.len(), 12); // 7 + 5
1928
1929        // Check that the run ends are correct
1930        let run_ends_values = result_run_array.run_ends().values();
1931        assert_eq!(&[2, 4, 7, 10, 12], run_ends_values);
1932
1933        // Check that the values are correct
1934        let expected = Int32Array::from(vec![Some(10), None, Some(30), Some(30), None]);
1935        let actual = result_run_array
1936            .values()
1937            .as_any()
1938            .downcast_ref::<Int32Array>()
1939            .unwrap();
1940        assert_eq!(actual.len(), expected.len());
1941        assert_eq!(actual.null_count(), expected.null_count());
1942        assert_eq!(actual.values(), expected.values());
1943    }
1944
1945    #[test]
1946    fn test_concat_run_array_single() {
1947        // Create a run array with run ends [2, 4] and values [10, 20]
1948        let run_ends1 = Int32Array::from(vec![2, 4]);
1949        let values1 = Int32Array::from(vec![10, 20]);
1950        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1951
1952        // Concatenate the single array
1953        let result = concat(&[&array1]).unwrap();
1954        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1955
1956        // The result should have length 4
1957        assert_eq!(result_run_array.len(), 4);
1958
1959        // Check that the run ends are correct
1960        let run_ends = result_run_array.run_ends().values();
1961        assert_eq!(&[2, 4], run_ends);
1962
1963        // Check that the values are correct
1964        assert_eq!(
1965            &[10, 20],
1966            result_run_array
1967                .values()
1968                .as_any()
1969                .downcast_ref::<Int32Array>()
1970                .unwrap()
1971                .values()
1972        );
1973    }
1974
1975    #[test]
1976    fn test_concat_run_array_with_3_arrays() {
1977        let run_ends1 = Int32Array::from(vec![2, 4]);
1978        let values1 = Int32Array::from(vec![10, 20]);
1979        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1980        let run_ends2 = Int32Array::from(vec![1, 4]);
1981        let values2 = Int32Array::from(vec![30, 40]);
1982        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
1983        let run_ends3 = Int32Array::from(vec![1, 4]);
1984        let values3 = Int32Array::from(vec![50, 60]);
1985        let array3 = RunArray::try_new(&run_ends3, &values3).unwrap();
1986
1987        // Concatenate the arrays
1988        let result = concat(&[&array1, &array2, &array3]).unwrap();
1989        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1990
1991        // Check that the result has the correct length
1992        assert_eq!(result_run_array.len(), 12); // 4 + 4 + 4
1993
1994        // Check the run ends
1995        let run_ends = result_run_array.run_ends().values();
1996        assert_eq!(run_ends.len(), 6);
1997        assert_eq!(&[2, 4, 5, 8, 9, 12], run_ends);
1998
1999        // Check the values
2000        let values = result_run_array
2001            .values()
2002            .as_any()
2003            .downcast_ref::<Int32Array>()
2004            .unwrap();
2005        assert_eq!(values.len(), 6);
2006        assert_eq!(&[10, 20, 30, 40, 50, 60], values.values());
2007    }
2008
2009    #[test]
2010    fn test_concat_run_array_with_truncated_run() {
2011        // Create a run array with run ends [2, 5] and values [10, 20]
2012        // Logical: [10, 10, 20, 20, 20]
2013        let run_ends1 = Int32Array::from(vec![2, 5]);
2014        let values1 = Int32Array::from(vec![10, 20]);
2015        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
2016        let array1_sliced = array1.slice(0, 3);
2017
2018        let run_ends2 = Int32Array::from(vec![2]);
2019        let values2 = Int32Array::from(vec![30]);
2020        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
2021
2022        let result = concat(&[&array1_sliced, &array2]).unwrap();
2023        let result_run_array = result.as_run::<Int32Type>();
2024
2025        // Result should be [10, 10, 20, 30, 30]
2026        // Run ends should be [2, 3, 5]
2027        assert_eq!(result_run_array.len(), 5);
2028        let run_ends = result_run_array.run_ends().values();
2029        let values = result_run_array.values().as_primitive::<Int32Type>();
2030        assert_eq!(values.values(), &[10, 20, 30]);
2031        assert_eq!(&[2, 3, 5], run_ends);
2032    }
2033
2034    /// A single row of a {String -> Int32} map: `None` for a null row, otherwise
2035    /// the list of (key, optional value) entries.
2036    type StringIntMapRow<'a> = Option<Vec<(&'a str, Option<i32>)>>;
2037
2038    /// Helper to build a MapArray of {String -> Int32} from a list of entries per row.
2039    fn build_string_int_map(rows: Vec<StringIntMapRow>) -> MapArray {
2040        let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32ArrayBuilder::new());
2041        for row in rows {
2042            match row {
2043                Some(entries) => {
2044                    for (k, v) in entries {
2045                        builder.keys().append_value(k);
2046                        builder.values().append_option(v);
2047                    }
2048                    builder.append(true).unwrap();
2049                }
2050                None => {
2051                    builder.append(false).unwrap();
2052                }
2053            }
2054        }
2055        builder.finish()
2056    }
2057
2058    #[test]
2059    fn test_concat_map_arrays() {
2060        let map1 = build_string_int_map(vec![
2061            Some(vec![("a", Some(1)), ("b", Some(2))]),
2062            Some(vec![("c", Some(3))]),
2063        ]);
2064        let map2 = build_string_int_map(vec![
2065            Some(vec![("d", Some(4)), ("e", Some(5))]),
2066            None,
2067            Some(vec![("f", Some(6))]),
2068        ]);
2069
2070        let result = concat(&[&map1, &map2]).unwrap();
2071        let result_map = result.as_map();
2072
2073        assert_eq!(result_map.len(), 5);
2074        assert_eq!(result_map.null_count(), 1);
2075
2076        // Check offsets
2077        assert_eq!(result_map.value_offsets(), &[0, 2, 3, 5, 5, 6]);
2078
2079        // Check keys
2080        let keys = result_map.keys().as_string::<i32>();
2081        let expected_keys: Vec<&str> = vec!["a", "b", "c", "d", "e", "f"];
2082        let actual_keys: Vec<&str> = keys.iter().map(|v| v.unwrap()).collect();
2083        assert_eq!(actual_keys, expected_keys);
2084
2085        // Check values
2086        let values = result_map.values().as_primitive::<Int32Type>();
2087        assert_eq!(values.values(), &[1, 2, 3, 4, 5, 6]);
2088    }
2089
2090    #[test]
2091    fn test_concat_map_arrays_sliced() {
2092        let map = build_string_int_map(vec![
2093            Some(vec![("a", Some(1))]),
2094            Some(vec![("b", Some(2)), ("c", Some(3))]),
2095            Some(vec![("d", Some(4))]),
2096            Some(vec![("e", Some(5))]),
2097        ]);
2098
2099        // Slice to get the middle two rows: [("b",2),("c",3)] and [("d",4)]
2100        let sliced = map.slice(1, 2);
2101
2102        let map2 = build_string_int_map(vec![Some(vec![("f", Some(6))])]);
2103
2104        let result = concat(&[&sliced, &map2]).unwrap();
2105        let result_map = result.as_map();
2106
2107        assert_eq!(result_map.len(), 3);
2108        assert_eq!(result_map.value_offsets(), &[0, 2, 3, 4]);
2109
2110        let keys = result_map.keys().as_string::<i32>();
2111        let actual_keys: Vec<&str> = keys.iter().map(|v| v.unwrap()).collect();
2112        assert_eq!(actual_keys, vec!["b", "c", "d", "f"]);
2113    }
2114
2115    #[test]
2116    fn test_concat_map_arrays_with_nulls() {
2117        let map1 = build_string_int_map(vec![Some(vec![("a", Some(1))]), None]);
2118        let map2 = build_string_int_map(vec![None, Some(vec![("b", Some(2))])]);
2119
2120        let result = concat(&[&map1, &map2]).unwrap();
2121        let result_map = result.as_map();
2122
2123        assert_eq!(result_map.len(), 4);
2124        assert_eq!(result_map.null_count(), 2);
2125        assert!(result_map.is_valid(0));
2126        assert!(result_map.is_null(1));
2127        assert!(result_map.is_null(2));
2128        assert!(result_map.is_valid(3));
2129    }
2130
2131    #[test]
2132    fn test_concat_map_arrays_empty_maps() {
2133        let map1 = build_string_int_map(vec![Some(vec![]), Some(vec![("a", Some(1))])]);
2134        let map2 = build_string_int_map(vec![
2135            Some(vec![]),
2136            Some(vec![("b", Some(2)), ("c", Some(3))]),
2137        ]);
2138
2139        let result = concat(&[&map1, &map2]).unwrap();
2140        let result_map = result.as_map();
2141
2142        assert_eq!(result_map.len(), 4);
2143        assert_eq!(result_map.null_count(), 0);
2144        assert_eq!(result_map.value_offsets(), &[0, 0, 1, 1, 3]);
2145    }
2146}