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    fn test_concat_13_incompatible_datatypes_should_not_include_all_of_them() {
803        let re = concat(&[
804            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(2), None]),
805            // 2 string to make sure we only mention unique types
806            &StringArray::from(vec![Some("hello"), Some("bar"), Some("world")]),
807            &StringArray::from(vec![Some("hey"), Some(""), Some("you")]),
808            // Another type to make sure we are showing all the incompatible types
809            &PrimitiveArray::<Int32Type>::from(vec![Some(-1), Some(2), None]),
810            &PrimitiveArray::<Int8Type>::from(vec![Some(-1), Some(2), None]),
811            &PrimitiveArray::<Int16Type>::from(vec![Some(-1), Some(2), None]),
812            &PrimitiveArray::<UInt8Type>::from(vec![Some(1), Some(2), None]),
813            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
814            &PrimitiveArray::<UInt32Type>::from(vec![Some(1), Some(2), None]),
815            // Non unique
816            &PrimitiveArray::<UInt16Type>::from(vec![Some(1), Some(2), None]),
817            &PrimitiveArray::<UInt64Type>::from(vec![Some(1), Some(2), None]),
818            &PrimitiveArray::<Float32Type>::from(vec![Some(1.0), Some(2.0), None]),
819            &PrimitiveArray::<Float64Type>::from(vec![Some(1.0), Some(2.0), None]),
820            &PrimitiveArray::<Float16Type>::new_null(3),
821            &BooleanArray::from(vec![Some(true), Some(false), None]),
822        ]);
823
824        assert_eq!(
825            re.unwrap_err().to_string(),
826            "Invalid argument error: It is not possible to concatenate arrays of different data types (Int64, Utf8, Int32, Int8, Int16, UInt8, UInt16, UInt32, UInt64, Float32, ...)."
827        );
828    }
829
830    #[test]
831    fn test_concat_string_arrays() {
832        let arr = concat(&[
833            &StringArray::from(vec!["hello", "world"]),
834            &StringArray::from(vec!["2", "3", "4"]),
835            &StringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]),
836        ])
837        .unwrap();
838
839        let expected_output = Arc::new(StringArray::from(vec![
840            Some("hello"),
841            Some("world"),
842            Some("2"),
843            Some("3"),
844            Some("4"),
845            Some("foo"),
846            Some("bar"),
847            None,
848            Some("baz"),
849        ])) as ArrayRef;
850
851        assert_eq!(&arr, &expected_output);
852    }
853
854    #[test]
855    fn test_concat_string_view_arrays() {
856        let arr = concat(&[
857            &StringViewArray::from(vec!["helloxxxxxxxxxxa", "world____________"]),
858            &StringViewArray::from(vec!["helloxxxxxxxxxxy", "3", "4"]),
859            &StringViewArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]),
860        ])
861        .unwrap();
862
863        let expected_output = Arc::new(StringViewArray::from(vec![
864            Some("helloxxxxxxxxxxa"),
865            Some("world____________"),
866            Some("helloxxxxxxxxxxy"),
867            Some("3"),
868            Some("4"),
869            Some("foo"),
870            Some("bar"),
871            None,
872            Some("baz"),
873        ])) as ArrayRef;
874
875        assert_eq!(&arr, &expected_output);
876    }
877
878    #[test]
879    fn test_concat_primitive_arrays() {
880        let arr = concat(&[
881            &PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(-1), Some(2), None, None]),
882            &PrimitiveArray::<Int64Type>::from(vec![Some(101), Some(102), Some(103), None]),
883            &PrimitiveArray::<Int64Type>::from(vec![Some(256), Some(512), Some(1024)]),
884        ])
885        .unwrap();
886
887        let expected_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
888            Some(-1),
889            Some(-1),
890            Some(2),
891            None,
892            None,
893            Some(101),
894            Some(102),
895            Some(103),
896            None,
897            Some(256),
898            Some(512),
899            Some(1024),
900        ])) as ArrayRef;
901
902        assert_eq!(&arr, &expected_output);
903    }
904
905    #[test]
906    fn test_concat_primitive_array_slices() {
907        let input_1 =
908            PrimitiveArray::<Int64Type>::from(vec![Some(-1), Some(-1), Some(2), None, None])
909                .slice(1, 3);
910
911        let input_2 =
912            PrimitiveArray::<Int64Type>::from(vec![Some(101), Some(102), Some(103), None])
913                .slice(1, 3);
914        let arr = concat(&[&input_1, &input_2]).unwrap();
915
916        let expected_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
917            Some(-1),
918            Some(2),
919            None,
920            Some(102),
921            Some(103),
922            None,
923        ])) as ArrayRef;
924
925        assert_eq!(&arr, &expected_output);
926    }
927
928    #[test]
929    fn test_concat_boolean_primitive_arrays() {
930        let arr = concat(&[
931            &BooleanArray::from(vec![
932                Some(true),
933                Some(true),
934                Some(false),
935                None,
936                None,
937                Some(false),
938            ]),
939            &BooleanArray::from(vec![None, Some(false), Some(true), Some(false)]),
940        ])
941        .unwrap();
942
943        let expected_output = Arc::new(BooleanArray::from(vec![
944            Some(true),
945            Some(true),
946            Some(false),
947            None,
948            None,
949            Some(false),
950            None,
951            Some(false),
952            Some(true),
953            Some(false),
954        ])) as ArrayRef;
955
956        assert_eq!(&arr, &expected_output);
957    }
958
959    #[test]
960    fn test_concat_primitive_list_arrays() {
961        let list1 = [
962            Some(vec![Some(-1), Some(-1), Some(2), None, None]),
963            Some(vec![]),
964            None,
965            Some(vec![Some(10)]),
966        ];
967        let list1_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone());
968
969        let list2 = [
970            None,
971            Some(vec![Some(100), None, Some(101)]),
972            Some(vec![Some(102)]),
973        ];
974        let list2_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone());
975
976        let list3 = [Some(vec![Some(1000), Some(1001)])];
977        let list3_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list3.clone());
978
979        let array_result = concat(&[&list1_array, &list2_array, &list3_array]).unwrap();
980
981        let expected = list1.into_iter().chain(list2).chain(list3);
982        let array_expected = ListArray::from_iter_primitive::<Int64Type, _, _>(expected);
983
984        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
985    }
986
987    #[test]
988    fn test_concat_primitive_list_arrays_slices() {
989        let list1 = [
990            Some(vec![Some(-1), Some(-1), Some(2), None, None]),
991            Some(vec![]), // In slice
992            None,         // In slice
993            Some(vec![Some(10)]),
994        ];
995        let list1_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone());
996        let list1_array = list1_array.slice(1, 2);
997        let list1_values = list1.into_iter().skip(1).take(2);
998
999        let list2 = [
1000            None,
1001            Some(vec![Some(100), None, Some(101)]),
1002            Some(vec![Some(102)]),
1003        ];
1004        let list2_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone());
1005
1006        // verify that this test covers the case when the first offset is non zero
1007        assert!(list1_array.offsets()[0].as_usize() > 0);
1008        let array_result = concat(&[&list1_array, &list2_array]).unwrap();
1009
1010        let expected = list1_values.chain(list2);
1011        let array_expected = ListArray::from_iter_primitive::<Int64Type, _, _>(expected);
1012
1013        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1014    }
1015
1016    #[test]
1017    fn test_concat_primitive_list_arrays_sliced_lengths() {
1018        let list1 = [
1019            Some(vec![Some(-1), Some(-1), Some(2), None, None]), // In slice
1020            Some(vec![]),                                        // In slice
1021            None,                                                // In slice
1022            Some(vec![Some(10)]),
1023        ];
1024        let list1_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone());
1025        let list1_array = list1_array.slice(0, 3); // no offset, but not all values
1026        let list1_values = list1.into_iter().take(3);
1027
1028        let list2 = [
1029            None,
1030            Some(vec![Some(100), None, Some(101)]),
1031            Some(vec![Some(102)]),
1032        ];
1033        let list2_array = ListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone());
1034
1035        // verify that this test covers the case when the first offset is zero, but the
1036        // last offset doesn't cover the entire array
1037        assert_eq!(list1_array.offsets()[0].as_usize(), 0);
1038        assert!(list1_array.offsets().last().unwrap().as_usize() < list1_array.values().len());
1039        let array_result = concat(&[&list1_array, &list2_array]).unwrap();
1040
1041        let expected = list1_values.chain(list2);
1042        let array_expected = ListArray::from_iter_primitive::<Int64Type, _, _>(expected);
1043
1044        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1045    }
1046
1047    #[test]
1048    fn test_concat_primitive_fixed_size_list_arrays() {
1049        let list1 = [
1050            Some(vec![Some(-1), None]),
1051            None,
1052            Some(vec![Some(10), Some(20)]),
1053        ];
1054        let list1_array =
1055            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(list1.clone(), 2);
1056
1057        let list2 = [
1058            None,
1059            Some(vec![Some(100), None]),
1060            Some(vec![Some(102), Some(103)]),
1061        ];
1062        let list2_array =
1063            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(list2.clone(), 2);
1064
1065        let list3 = [Some(vec![Some(1000), Some(1001)])];
1066        let list3_array =
1067            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(list3.clone(), 2);
1068
1069        let array_result = concat(&[&list1_array, &list2_array, &list3_array]).unwrap();
1070
1071        let expected = list1.into_iter().chain(list2).chain(list3);
1072        let array_expected =
1073            FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(expected, 2);
1074
1075        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1076    }
1077
1078    #[test]
1079    fn test_concat_list_view_arrays() {
1080        let list1 = [
1081            Some(vec![Some(-1), None]),
1082            None,
1083            Some(vec![Some(10), Some(20)]),
1084        ];
1085        let mut list1_array = ListViewBuilder::new(Int64Builder::new());
1086        for v in list1.iter() {
1087            list1_array.append_option(v.clone());
1088        }
1089        let list1_array = list1_array.finish();
1090
1091        let list2 = [
1092            None,
1093            Some(vec![Some(100), None]),
1094            Some(vec![Some(102), Some(103)]),
1095        ];
1096        let mut list2_array = ListViewBuilder::new(Int64Builder::new());
1097        for v in list2.iter() {
1098            list2_array.append_option(v.clone());
1099        }
1100        let list2_array = list2_array.finish();
1101
1102        let list3 = [Some(vec![Some(1000), Some(1001)])];
1103        let mut list3_array = ListViewBuilder::new(Int64Builder::new());
1104        for v in list3.iter() {
1105            list3_array.append_option(v.clone());
1106        }
1107        let list3_array = list3_array.finish();
1108
1109        let array_result = concat(&[&list1_array, &list2_array, &list3_array]).unwrap();
1110
1111        let expected: Vec<_> = list1.into_iter().chain(list2).chain(list3).collect();
1112        let mut array_expected = ListViewBuilder::new(Int64Builder::new());
1113        for v in expected.iter() {
1114            array_expected.append_option(v.clone());
1115        }
1116        let array_expected = array_expected.finish();
1117
1118        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1119    }
1120
1121    #[test]
1122    fn test_concat_sliced_list_view_arrays() {
1123        let list1 = [
1124            Some(vec![Some(-1), None]),
1125            None,
1126            Some(vec![Some(10), Some(20)]),
1127        ];
1128        let mut list1_array = ListViewBuilder::new(Int64Builder::new());
1129        for v in list1.iter() {
1130            list1_array.append_option(v.clone());
1131        }
1132        let list1_array = list1_array.finish();
1133
1134        let list2 = [
1135            None,
1136            Some(vec![Some(100), None]),
1137            Some(vec![Some(102), Some(103)]),
1138        ];
1139        let mut list2_array = ListViewBuilder::new(Int64Builder::new());
1140        for v in list2.iter() {
1141            list2_array.append_option(v.clone());
1142        }
1143        let list2_array = list2_array.finish();
1144
1145        let list3 = [Some(vec![Some(1000), Some(1001)])];
1146        let mut list3_array = ListViewBuilder::new(Int64Builder::new());
1147        for v in list3.iter() {
1148            list3_array.append_option(v.clone());
1149        }
1150        let list3_array = list3_array.finish();
1151
1152        // Concat sliced arrays.
1153        // ListView slicing will slice the offset/sizes but preserve the original values child.
1154        let array_result = concat(&[
1155            &list1_array.slice(1, 2),
1156            &list2_array.slice(1, 2),
1157            &list3_array.slice(0, 1),
1158        ])
1159        .unwrap();
1160
1161        let expected: Vec<_> = vec![
1162            None,
1163            Some(vec![Some(10), Some(20)]),
1164            Some(vec![Some(100), None]),
1165            Some(vec![Some(102), Some(103)]),
1166            Some(vec![Some(1000), Some(1001)]),
1167        ];
1168        let mut array_expected = ListViewBuilder::new(Int64Builder::new());
1169        for v in expected.iter() {
1170            array_expected.append_option(v.clone());
1171        }
1172        let array_expected = array_expected.finish();
1173
1174        assert_eq!(array_result.as_ref(), &array_expected as &dyn Array);
1175    }
1176
1177    #[test]
1178    fn test_concat_struct_arrays() {
1179        let field = Arc::new(Field::new("field", DataType::Int64, true));
1180        let input_primitive_1: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1181            Some(-1),
1182            Some(-1),
1183            Some(2),
1184            None,
1185            None,
1186        ]));
1187        let input_struct_1 = StructArray::from(vec![(field.clone(), input_primitive_1)]);
1188
1189        let input_primitive_2: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1190            Some(101),
1191            Some(102),
1192            Some(103),
1193            None,
1194        ]));
1195        let input_struct_2 = StructArray::from(vec![(field.clone(), input_primitive_2)]);
1196
1197        let input_primitive_3: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1198            Some(256),
1199            Some(512),
1200            Some(1024),
1201        ]));
1202        let input_struct_3 = StructArray::from(vec![(field, input_primitive_3)]);
1203
1204        let arr = concat(&[&input_struct_1, &input_struct_2, &input_struct_3]).unwrap();
1205
1206        let expected_primitive_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1207            Some(-1),
1208            Some(-1),
1209            Some(2),
1210            None,
1211            None,
1212            Some(101),
1213            Some(102),
1214            Some(103),
1215            None,
1216            Some(256),
1217            Some(512),
1218            Some(1024),
1219        ])) as ArrayRef;
1220
1221        let actual_primitive = arr
1222            .as_any()
1223            .downcast_ref::<StructArray>()
1224            .unwrap()
1225            .column(0);
1226        assert_eq!(actual_primitive, &expected_primitive_output);
1227    }
1228
1229    #[test]
1230    fn test_concat_struct_array_slices() {
1231        let field = Arc::new(Field::new("field", DataType::Int64, true));
1232        let input_primitive_1: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1233            Some(-1),
1234            Some(-1),
1235            Some(2),
1236            None,
1237            None,
1238        ]));
1239        let input_struct_1 = StructArray::from(vec![(field.clone(), input_primitive_1)]);
1240
1241        let input_primitive_2: ArrayRef = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1242            Some(101),
1243            Some(102),
1244            Some(103),
1245            None,
1246        ]));
1247        let input_struct_2 = StructArray::from(vec![(field, input_primitive_2)]);
1248
1249        let arr = concat(&[&input_struct_1.slice(1, 3), &input_struct_2.slice(1, 2)]).unwrap();
1250
1251        let expected_primitive_output = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
1252            Some(-1),
1253            Some(2),
1254            None,
1255            Some(102),
1256            Some(103),
1257        ])) as ArrayRef;
1258
1259        let actual_primitive = arr
1260            .as_any()
1261            .downcast_ref::<StructArray>()
1262            .unwrap()
1263            .column(0);
1264        assert_eq!(actual_primitive, &expected_primitive_output);
1265    }
1266
1267    #[test]
1268    fn test_concat_struct_arrays_no_nulls() {
1269        let input_1a = vec![1, 2, 3];
1270        let input_1b = vec!["one", "two", "three"];
1271        let input_2a = vec![4, 5, 6, 7];
1272        let input_2b = vec!["four", "five", "six", "seven"];
1273
1274        let struct_from_primitives = |ints: Vec<i64>, strings: Vec<&str>| {
1275            StructArray::try_from(vec![
1276                ("ints", Arc::new(Int64Array::from(ints)) as _),
1277                ("strings", Arc::new(StringArray::from(strings)) as _),
1278            ])
1279        };
1280
1281        let expected_output = struct_from_primitives(
1282            [input_1a.clone(), input_2a.clone()].concat(),
1283            [input_1b.clone(), input_2b.clone()].concat(),
1284        )
1285        .unwrap();
1286
1287        let input_1 = struct_from_primitives(input_1a, input_1b).unwrap();
1288        let input_2 = struct_from_primitives(input_2a, input_2b).unwrap();
1289
1290        let arr = concat(&[&input_1, &input_2]).unwrap();
1291        let struct_result = arr.as_struct();
1292
1293        assert_eq!(struct_result, &expected_output);
1294        assert_eq!(arr.null_count(), 0);
1295    }
1296
1297    #[test]
1298    fn test_concat_struct_no_fields() {
1299        let input_1 = StructArray::new_empty_fields(10, None);
1300        let input_2 = StructArray::new_empty_fields(10, None);
1301        let arr = concat(&[&input_1, &input_2]).unwrap();
1302
1303        assert_eq!(arr.len(), 20);
1304        assert_eq!(arr.null_count(), 0);
1305
1306        let input1_valid = StructArray::new_empty_fields(10, Some(NullBuffer::new_valid(10)));
1307        let input2_null = StructArray::new_empty_fields(10, Some(NullBuffer::new_null(10)));
1308        let arr = concat(&[&input1_valid, &input2_null]).unwrap();
1309
1310        assert_eq!(arr.len(), 20);
1311        assert_eq!(arr.null_count(), 10);
1312    }
1313
1314    #[test]
1315    fn test_string_array_slices() {
1316        let input_1 = StringArray::from(vec!["hello", "A", "B", "C"]);
1317        let input_2 = StringArray::from(vec!["world", "D", "E", "Z"]);
1318
1319        let arr = concat(&[&input_1.slice(1, 3), &input_2.slice(1, 2)]).unwrap();
1320
1321        let expected_output = StringArray::from(vec!["A", "B", "C", "D", "E"]);
1322
1323        let actual_output = arr.as_any().downcast_ref::<StringArray>().unwrap();
1324        assert_eq!(actual_output, &expected_output);
1325    }
1326
1327    #[test]
1328    fn test_string_array_with_null_slices() {
1329        let input_1 = StringArray::from(vec![Some("hello"), None, Some("A"), Some("C")]);
1330        let input_2 = StringArray::from(vec![None, Some("world"), Some("D"), None]);
1331
1332        let arr = concat(&[&input_1.slice(1, 3), &input_2.slice(1, 2)]).unwrap();
1333
1334        let expected_output =
1335            StringArray::from(vec![None, Some("A"), Some("C"), Some("world"), Some("D")]);
1336
1337        let actual_output = arr.as_any().downcast_ref::<StringArray>().unwrap();
1338        assert_eq!(actual_output, &expected_output);
1339    }
1340
1341    fn collect_string_dictionary(array: &DictionaryArray<Int32Type>) -> Vec<Option<&str>> {
1342        let concrete = array.downcast_dict::<StringArray>().unwrap();
1343        concrete.into_iter().collect()
1344    }
1345
1346    #[test]
1347    fn test_string_dictionary_array() {
1348        let input_1: DictionaryArray<Int32Type> = vec!["hello", "A", "B", "hello", "hello", "C"]
1349            .into_iter()
1350            .collect();
1351        let input_2: DictionaryArray<Int32Type> = vec!["hello", "E", "E", "hello", "F", "E"]
1352            .into_iter()
1353            .collect();
1354
1355        let expected: Vec<_> = vec![
1356            "hello", "A", "B", "hello", "hello", "C", "hello", "E", "E", "hello", "F", "E",
1357        ]
1358        .into_iter()
1359        .map(Some)
1360        .collect();
1361
1362        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1363        let dictionary = concat.as_dictionary::<Int32Type>();
1364        let actual = collect_string_dictionary(dictionary);
1365        assert_eq!(actual, expected);
1366
1367        // Should have concatenated inputs together
1368        assert_eq!(
1369            dictionary.values().len(),
1370            input_1.values().len() + input_2.values().len(),
1371        )
1372    }
1373
1374    #[test]
1375    fn test_string_dictionary_array_nulls() {
1376        let input_1: DictionaryArray<Int32Type> = vec![Some("foo"), Some("bar"), None, Some("fiz")]
1377            .into_iter()
1378            .collect();
1379        let input_2: DictionaryArray<Int32Type> = vec![None].into_iter().collect();
1380        let expected = vec![Some("foo"), Some("bar"), None, Some("fiz"), None];
1381
1382        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1383        let dictionary = concat.as_dictionary::<Int32Type>();
1384        let actual = collect_string_dictionary(dictionary);
1385        assert_eq!(actual, expected);
1386
1387        // Should have concatenated inputs together
1388        assert_eq!(
1389            dictionary.values().len(),
1390            input_1.values().len() + input_2.values().len(),
1391        )
1392    }
1393
1394    #[test]
1395    fn test_string_dictionary_array_nulls_in_values() {
1396        let input_1_keys = Int32Array::from_iter_values([0, 2, 1, 3]);
1397        let input_1_values = StringArray::from(vec![Some("foo"), None, Some("bar"), Some("fiz")]);
1398        let input_1 = DictionaryArray::new(input_1_keys, Arc::new(input_1_values));
1399
1400        let input_2_keys = Int32Array::from_iter_values([0]);
1401        let input_2_values = StringArray::from(vec![None, Some("hello")]);
1402        let input_2 = DictionaryArray::new(input_2_keys, Arc::new(input_2_values));
1403
1404        let expected = vec![Some("foo"), Some("bar"), None, Some("fiz"), None];
1405
1406        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1407        let dictionary = concat.as_dictionary::<Int32Type>();
1408        let actual = collect_string_dictionary(dictionary);
1409        assert_eq!(actual, expected);
1410    }
1411
1412    #[test]
1413    fn test_string_dictionary_merge() {
1414        let mut builder = StringDictionaryBuilder::<Int32Type>::new();
1415        for i in 0..20 {
1416            builder.append(i.to_string()).unwrap();
1417        }
1418        let input_1 = builder.finish();
1419
1420        let mut builder = StringDictionaryBuilder::<Int32Type>::new();
1421        for i in 0..30 {
1422            builder.append(i.to_string()).unwrap();
1423        }
1424        let input_2 = builder.finish();
1425
1426        let expected: Vec<_> = (0..20).chain(0..30).map(|x| x.to_string()).collect();
1427        let expected: Vec<_> = expected.iter().map(|x| Some(x.as_str())).collect();
1428
1429        let concat = concat(&[&input_1 as _, &input_2 as _]).unwrap();
1430        let dictionary = concat.as_dictionary::<Int32Type>();
1431        let actual = collect_string_dictionary(dictionary);
1432        assert_eq!(actual, expected);
1433
1434        // Should have merged inputs together
1435        // Not 30 as this is done on a best-effort basis
1436        let values_len = dictionary.values().len();
1437        assert!((30..40).contains(&values_len), "{values_len}")
1438    }
1439
1440    #[test]
1441    fn test_primitive_dictionary_merge() {
1442        // Same value repeated 5 times.
1443        let keys = vec![1; 5];
1444        let values = (10..20).collect::<Vec<_>>();
1445        let dict = DictionaryArray::new(
1446            Int8Array::from(keys.clone()),
1447            Arc::new(Int32Array::from(values.clone())),
1448        );
1449        let other = DictionaryArray::new(
1450            Int8Array::from(keys.clone()),
1451            Arc::new(Int32Array::from(values.clone())),
1452        );
1453
1454        let result_same_dictionary = concat(&[&dict, &dict]).unwrap();
1455        // Verify pointer equality check succeeds, and therefore the
1456        // dictionaries are not merged. A single values buffer should be reused
1457        // in this case.
1458        assert!(
1459            dict.values().to_data().ptr_eq(
1460                &result_same_dictionary
1461                    .as_dictionary::<Int8Type>()
1462                    .values()
1463                    .to_data()
1464            )
1465        );
1466        assert_eq!(
1467            result_same_dictionary
1468                .as_dictionary::<Int8Type>()
1469                .values()
1470                .len(),
1471            values.len(),
1472        );
1473
1474        let result_cloned_dictionary = concat(&[&dict, &other]).unwrap();
1475        // Should have only 1 underlying value since all keys reference it.
1476        assert_eq!(
1477            result_cloned_dictionary
1478                .as_dictionary::<Int8Type>()
1479                .values()
1480                .len(),
1481            1
1482        );
1483    }
1484
1485    #[test]
1486    fn test_concat_string_sizes() {
1487        let a: LargeStringArray = ((0..150).map(|_| Some("foo"))).collect();
1488        let b: LargeStringArray = ((0..150).map(|_| Some("foo"))).collect();
1489        let c = LargeStringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]);
1490        // 150 * 3 = 450
1491        // 150 * 3 = 450
1492        // 3 * 3   = 9
1493        // ------------+
1494        // 909
1495
1496        let arr = concat(&[&a, &b, &c]).unwrap();
1497        assert_eq!(arr.to_data().buffers()[1].capacity(), 909);
1498    }
1499
1500    #[test]
1501    fn test_dictionary_concat_reuse() {
1502        let array: DictionaryArray<Int8Type> = vec!["a", "a", "b", "c"].into_iter().collect();
1503        let copy: DictionaryArray<Int8Type> = array.clone();
1504
1505        // dictionary is "a", "b", "c"
1506        assert_eq!(
1507            array.values(),
1508            &(Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef)
1509        );
1510        assert_eq!(array.keys(), &Int8Array::from(vec![0, 0, 1, 2]));
1511
1512        // concatenate it with itself
1513        let combined = concat(&[&copy as _, &array as _]).unwrap();
1514        let combined = combined.as_dictionary::<Int8Type>();
1515
1516        assert_eq!(
1517            combined.values(),
1518            &(Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef),
1519            "Actual: {combined:#?}"
1520        );
1521
1522        assert_eq!(
1523            combined.keys(),
1524            &Int8Array::from(vec![0, 0, 1, 2, 0, 0, 1, 2])
1525        );
1526
1527        // Should have reused the dictionary
1528        assert!(
1529            array
1530                .values()
1531                .to_data()
1532                .ptr_eq(&combined.values().to_data())
1533        );
1534        assert!(copy.values().to_data().ptr_eq(&combined.values().to_data()));
1535
1536        let new: DictionaryArray<Int8Type> = vec!["d"].into_iter().collect();
1537        let combined = concat(&[&copy as _, &array as _, &new as _]).unwrap();
1538        let com = combined.as_dictionary::<Int8Type>();
1539
1540        // Should not have reused the dictionary
1541        assert!(!array.values().to_data().ptr_eq(&com.values().to_data()));
1542        assert!(!copy.values().to_data().ptr_eq(&com.values().to_data()));
1543        assert!(!new.values().to_data().ptr_eq(&com.values().to_data()));
1544    }
1545
1546    #[test]
1547    fn concat_record_batches() {
1548        let schema = Arc::new(Schema::new(vec![
1549            Field::new("a", DataType::Int32, false),
1550            Field::new("b", DataType::Utf8, false),
1551        ]));
1552        let batch1 = RecordBatch::try_new(
1553            schema.clone(),
1554            vec![
1555                Arc::new(Int32Array::from(vec![1, 2])),
1556                Arc::new(StringArray::from(vec!["a", "b"])),
1557            ],
1558        )
1559        .unwrap();
1560        let batch2 = RecordBatch::try_new(
1561            schema.clone(),
1562            vec![
1563                Arc::new(Int32Array::from(vec![3, 4])),
1564                Arc::new(StringArray::from(vec!["c", "d"])),
1565            ],
1566        )
1567        .unwrap();
1568        let new_batch = concat_batches(&schema, [&batch1, &batch2]).unwrap();
1569        assert_eq!(new_batch.schema().as_ref(), schema.as_ref());
1570        assert_eq!(2, new_batch.num_columns());
1571        assert_eq!(4, new_batch.num_rows());
1572        let new_batch_owned = concat_batches(&schema, &[batch1, batch2]).unwrap();
1573        assert_eq!(new_batch_owned.schema().as_ref(), schema.as_ref());
1574        assert_eq!(2, new_batch_owned.num_columns());
1575        assert_eq!(4, new_batch_owned.num_rows());
1576    }
1577
1578    #[test]
1579    fn concat_empty_record_batch() {
1580        let schema = Arc::new(Schema::new(vec![
1581            Field::new("a", DataType::Int32, false),
1582            Field::new("b", DataType::Utf8, false),
1583        ]));
1584        let batch = concat_batches(&schema, []).unwrap();
1585        assert_eq!(batch.schema().as_ref(), schema.as_ref());
1586        assert_eq!(0, batch.num_rows());
1587    }
1588
1589    #[test]
1590    fn concat_record_batches_of_different_schemas_but_compatible_data() {
1591        let schema1 = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
1592        // column names differ
1593        let schema2 = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, false)]));
1594        let batch1 = RecordBatch::try_new(
1595            schema1.clone(),
1596            vec![Arc::new(Int32Array::from(vec![1, 2]))],
1597        )
1598        .unwrap();
1599        let batch2 =
1600            RecordBatch::try_new(schema2, vec![Arc::new(Int32Array::from(vec![3, 4]))]).unwrap();
1601        // concat_batches simply uses the schema provided
1602        let batch = concat_batches(&schema1, [&batch1, &batch2]).unwrap();
1603        assert_eq!(batch.schema().as_ref(), schema1.as_ref());
1604        assert_eq!(4, batch.num_rows());
1605    }
1606
1607    #[test]
1608    fn concat_record_batches_of_different_schemas_incompatible_data() {
1609        let schema1 = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
1610        // column names differ
1611        let schema2 = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, false)]));
1612        let batch1 = RecordBatch::try_new(
1613            schema1.clone(),
1614            vec![Arc::new(Int32Array::from(vec![1, 2]))],
1615        )
1616        .unwrap();
1617        let batch2 = RecordBatch::try_new(
1618            schema2,
1619            vec![Arc::new(StringArray::from(vec!["foo", "bar"]))],
1620        )
1621        .unwrap();
1622
1623        let error = concat_batches(&schema1, [&batch1, &batch2]).unwrap_err();
1624        assert_eq!(
1625            error.to_string(),
1626            "Invalid argument error: It is not possible to concatenate arrays of different data types (Int32, Utf8)."
1627        );
1628    }
1629
1630    #[test]
1631    fn concat_capacity() {
1632        let a = Int32Array::from_iter_values(0..100);
1633        let b = Int32Array::from_iter_values(10..20);
1634        let a = concat(&[&a, &b]).unwrap();
1635        let data = a.to_data();
1636        assert_eq!(data.buffers()[0].len(), 440);
1637        assert_eq!(data.buffers()[0].capacity(), 440);
1638
1639        let a = concat(&[&a.slice(10, 20), &b]).unwrap();
1640        let data = a.to_data();
1641        assert_eq!(data.buffers()[0].len(), 120);
1642        assert_eq!(data.buffers()[0].capacity(), 120);
1643
1644        let a = StringArray::from_iter_values(std::iter::repeat_n("foo", 100));
1645        let b = StringArray::from(vec!["bingo", "bongo", "lorem", ""]);
1646
1647        let a = concat(&[&a, &b]).unwrap();
1648        let data = a.to_data();
1649        // (100 + 4 + 1) * size_of<i32>()
1650        assert_eq!(data.buffers()[0].len(), 420);
1651        assert_eq!(data.buffers()[0].capacity(), 420);
1652
1653        // len("foo") * 100 + len("bingo") + len("bongo") + len("lorem")
1654        assert_eq!(data.buffers()[1].len(), 315);
1655        assert_eq!(data.buffers()[1].capacity(), 315);
1656
1657        let a = concat(&[&a.slice(10, 40), &b]).unwrap();
1658        let data = a.to_data();
1659        // (40 + 4 + 5) * size_of<i32>()
1660        assert_eq!(data.buffers()[0].len(), 180);
1661        assert_eq!(data.buffers()[0].capacity(), 180);
1662
1663        // len("foo") * 40 + len("bingo") + len("bongo") + len("lorem")
1664        assert_eq!(data.buffers()[1].len(), 135);
1665        assert_eq!(data.buffers()[1].capacity(), 135);
1666
1667        let a = LargeBinaryArray::from_iter_values(std::iter::repeat_n(b"foo", 100));
1668        let b = LargeBinaryArray::from_iter_values(std::iter::repeat_n(b"cupcakes", 10));
1669
1670        let a = concat(&[&a, &b]).unwrap();
1671        let data = a.to_data();
1672        // (100 + 10 + 1) * size_of<i64>()
1673        assert_eq!(data.buffers()[0].len(), 888);
1674        assert_eq!(data.buffers()[0].capacity(), 888);
1675
1676        // len("foo") * 100 + len("cupcakes") * 10
1677        assert_eq!(data.buffers()[1].len(), 380);
1678        assert_eq!(data.buffers()[1].capacity(), 380);
1679
1680        let a = concat(&[&a.slice(10, 40), &b]).unwrap();
1681        let data = a.to_data();
1682        // (40 + 10 + 1) * size_of<i64>()
1683        assert_eq!(data.buffers()[0].len(), 408);
1684        assert_eq!(data.buffers()[0].capacity(), 408);
1685
1686        // len("foo") * 40 + len("cupcakes") * 10
1687        assert_eq!(data.buffers()[1].len(), 200);
1688        assert_eq!(data.buffers()[1].capacity(), 200);
1689    }
1690
1691    #[test]
1692    fn concat_sparse_nulls() {
1693        let values = StringArray::from_iter_values((0..100).map(|x| x.to_string()));
1694        let keys = Int32Array::from(vec![1; 10]);
1695        let dict_a = DictionaryArray::new(keys, Arc::new(values));
1696        let values = StringArray::new_null(0);
1697        let keys = Int32Array::new_null(10);
1698        let dict_b = DictionaryArray::new(keys, Arc::new(values));
1699        let array = concat(&[&dict_a, &dict_b]).unwrap();
1700        assert_eq!(array.null_count(), 10);
1701        assert_eq!(array.logical_null_count(), 10);
1702    }
1703
1704    #[test]
1705    fn concat_dictionary_list_array_simple() {
1706        let scalars = [
1707            create_single_row_list_of_dict(vec![Some("a")]),
1708            create_single_row_list_of_dict(vec![Some("a")]),
1709            create_single_row_list_of_dict(vec![Some("b")]),
1710        ];
1711
1712        let arrays = scalars.iter().map(|a| a as &dyn Array).collect::<Vec<_>>();
1713        let concat_res = concat(arrays.as_slice()).unwrap();
1714
1715        let expected_list = create_list_of_dict(vec![
1716            // Row 1
1717            Some(vec![Some("a")]),
1718            Some(vec![Some("a")]),
1719            Some(vec![Some("b")]),
1720        ]);
1721
1722        let list = concat_res.as_list::<i32>();
1723
1724        // Assert that the list is equal to the expected list
1725        list.iter().zip(expected_list.iter()).for_each(|(a, b)| {
1726            assert_eq!(a, b);
1727        });
1728
1729        assert_dictionary_has_unique_values::<_, StringArray>(
1730            list.values().as_dictionary::<Int32Type>(),
1731        );
1732    }
1733
1734    #[test]
1735    fn concat_many_dictionary_list_arrays() {
1736        let number_of_unique_values = 8;
1737        let scalars = (0..80000)
1738            .map(|i| {
1739                create_single_row_list_of_dict(vec![Some(
1740                    (i % number_of_unique_values).to_string(),
1741                )])
1742            })
1743            .collect::<Vec<_>>();
1744
1745        let arrays = scalars.iter().map(|a| a as &dyn Array).collect::<Vec<_>>();
1746        let concat_res = concat(arrays.as_slice()).unwrap();
1747
1748        let expected_list = create_list_of_dict(
1749            (0..80000)
1750                .map(|i| Some(vec![Some((i % number_of_unique_values).to_string())]))
1751                .collect::<Vec<_>>(),
1752        );
1753
1754        let list = concat_res.as_list::<i32>();
1755
1756        // Assert that the list is equal to the expected list
1757        list.iter().zip(expected_list.iter()).for_each(|(a, b)| {
1758            assert_eq!(a, b);
1759        });
1760
1761        assert_dictionary_has_unique_values::<_, StringArray>(
1762            list.values().as_dictionary::<Int32Type>(),
1763        );
1764    }
1765
1766    fn create_single_row_list_of_dict(
1767        list_items: Vec<Option<impl AsRef<str>>>,
1768    ) -> GenericListArray<i32> {
1769        let rows = list_items.into_iter().map(Some).collect();
1770
1771        create_list_of_dict(vec![rows])
1772    }
1773
1774    fn create_list_of_dict(
1775        rows: Vec<Option<Vec<Option<impl AsRef<str>>>>>,
1776    ) -> GenericListArray<i32> {
1777        let mut builder =
1778            GenericListBuilder::<i32, _>::new(StringDictionaryBuilder::<Int32Type>::new());
1779
1780        for row in rows {
1781            builder.append_option(row);
1782        }
1783
1784        builder.finish()
1785    }
1786
1787    fn assert_dictionary_has_unique_values<'a, K, V>(array: &'a DictionaryArray<K>)
1788    where
1789        K: ArrowDictionaryKeyType,
1790        V: Sync + Send + 'static,
1791        &'a V: ArrayAccessor + IntoIterator,
1792        <&'a V as ArrayAccessor>::Item: Default + Clone + PartialEq + Debug + Ord,
1793        <&'a V as IntoIterator>::Item: Clone + PartialEq + Debug + Ord,
1794    {
1795        let dict = array.downcast_dict::<V>().unwrap();
1796        let mut values = dict.values().into_iter().collect::<Vec<_>>();
1797
1798        // remove duplicates must be sorted first so we can compare
1799        values.sort();
1800
1801        let mut unique_values = values.clone();
1802
1803        unique_values.dedup();
1804
1805        assert_eq!(
1806            values, unique_values,
1807            "There are duplicates in the value list (the value list here is sorted which is only for the assertion)"
1808        );
1809    }
1810
1811    // Test the simple case of concatenating two RunArrays
1812    #[test]
1813    fn test_concat_run_array() {
1814        // Create simple run arrays
1815        let run_ends1 = Int32Array::from(vec![2, 4]);
1816        let values1 = Int32Array::from(vec![10, 20]);
1817        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1818
1819        let run_ends2 = Int32Array::from(vec![1, 4]);
1820        let values2 = Int32Array::from(vec![30, 40]);
1821        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
1822
1823        // Concatenate the arrays - this should now work properly
1824        let result = concat(&[&array1, &array2]).unwrap();
1825        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1826
1827        // Check that the result has the correct length
1828        assert_eq!(result_run_array.len(), 8); // 4 + 4
1829
1830        // Check the run ends
1831        let run_ends = result_run_array.run_ends().values();
1832        assert_eq!(run_ends.len(), 4);
1833        assert_eq!(&[2, 4, 5, 8], run_ends);
1834
1835        // Check the values
1836        let values = result_run_array
1837            .values()
1838            .as_any()
1839            .downcast_ref::<Int32Array>()
1840            .unwrap();
1841        assert_eq!(values.len(), 4);
1842        assert_eq!(&[10, 20, 30, 40], values.values());
1843    }
1844
1845    #[test]
1846    fn test_concat_sliced_run_array() {
1847        // Slicing away first run in both arrays
1848        let run_ends1 = Int32Array::from(vec![2, 4]);
1849        let values1 = Int32Array::from(vec![10, 20]);
1850        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap(); // [10, 10, 20, 20]
1851        let array1 = array1.slice(2, 2); // [20, 20]
1852
1853        let run_ends2 = Int32Array::from(vec![1, 4]);
1854        let values2 = Int32Array::from(vec![30, 40]);
1855        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap(); // [30, 40, 40, 40]
1856        let array2 = array2.slice(1, 3); // [40, 40, 40]
1857
1858        let result = concat(&[&array1, &array2]).unwrap();
1859        let result = result.as_run::<Int32Type>();
1860        let result = result.downcast::<Int32Array>().unwrap();
1861
1862        let expected = vec![20, 20, 40, 40, 40];
1863        let actual = result.into_iter().flatten().collect::<Vec<_>>();
1864        assert_eq!(expected, actual);
1865    }
1866
1867    #[test]
1868    fn test_concat_run_array_matching_first_last_value() {
1869        // Create a run array with run ends [2, 4, 7] and values [10, 20, 30]
1870        let run_ends1 = Int32Array::from(vec![2, 4, 7]);
1871        let values1 = Int32Array::from(vec![10, 20, 30]);
1872        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1873
1874        // Create another run array with run ends [3, 5] and values [30, 40]
1875        let run_ends2 = Int32Array::from(vec![3, 5]);
1876        let values2 = Int32Array::from(vec![30, 40]);
1877        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
1878
1879        // Concatenate the two arrays
1880        let result = concat(&[&array1, &array2]).unwrap();
1881        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1882
1883        // The result should have length 12 (7 + 5)
1884        assert_eq!(result_run_array.len(), 12);
1885
1886        // Check that the run ends are correct
1887        let run_ends = result_run_array.run_ends().values();
1888        assert_eq!(&[2, 4, 7, 10, 12], run_ends);
1889
1890        // Check that the values are correct
1891        assert_eq!(
1892            &[10, 20, 30, 30, 40],
1893            result_run_array
1894                .values()
1895                .as_any()
1896                .downcast_ref::<Int32Array>()
1897                .unwrap()
1898                .values()
1899        );
1900    }
1901
1902    #[test]
1903    fn test_concat_run_array_with_nulls() {
1904        // Create values array with nulls
1905        let values1 = Int32Array::from(vec![Some(10), None, Some(30)]);
1906        let run_ends1 = Int32Array::from(vec![2, 4, 7]);
1907        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1908
1909        // Create another run array with run ends [3, 5] and values [30, null]
1910        let values2 = Int32Array::from(vec![Some(30), None]);
1911        let run_ends2 = Int32Array::from(vec![3, 5]);
1912        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
1913
1914        // Concatenate the two arrays
1915        let result = concat(&[&array1, &array2]).unwrap();
1916        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1917
1918        // The result should have length 12 (7 + 5)
1919        assert_eq!(result_run_array.len(), 12);
1920
1921        // Get a reference to the run array itself for testing
1922
1923        // Just test the length and run ends without asserting specific values
1924        // This ensures the test passes while we work on full support for RunArray nulls
1925        assert_eq!(result_run_array.len(), 12); // 7 + 5
1926
1927        // Check that the run ends are correct
1928        let run_ends_values = result_run_array.run_ends().values();
1929        assert_eq!(&[2, 4, 7, 10, 12], run_ends_values);
1930
1931        // Check that the values are correct
1932        let expected = Int32Array::from(vec![Some(10), None, Some(30), Some(30), None]);
1933        let actual = result_run_array
1934            .values()
1935            .as_any()
1936            .downcast_ref::<Int32Array>()
1937            .unwrap();
1938        assert_eq!(actual.len(), expected.len());
1939        assert_eq!(actual.null_count(), expected.null_count());
1940        assert_eq!(actual.values(), expected.values());
1941    }
1942
1943    #[test]
1944    fn test_concat_run_array_single() {
1945        // Create a run array with run ends [2, 4] and values [10, 20]
1946        let run_ends1 = Int32Array::from(vec![2, 4]);
1947        let values1 = Int32Array::from(vec![10, 20]);
1948        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1949
1950        // Concatenate the single array
1951        let result = concat(&[&array1]).unwrap();
1952        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1953
1954        // The result should have length 4
1955        assert_eq!(result_run_array.len(), 4);
1956
1957        // Check that the run ends are correct
1958        let run_ends = result_run_array.run_ends().values();
1959        assert_eq!(&[2, 4], run_ends);
1960
1961        // Check that the values are correct
1962        assert_eq!(
1963            &[10, 20],
1964            result_run_array
1965                .values()
1966                .as_any()
1967                .downcast_ref::<Int32Array>()
1968                .unwrap()
1969                .values()
1970        );
1971    }
1972
1973    #[test]
1974    fn test_concat_run_array_with_3_arrays() {
1975        let run_ends1 = Int32Array::from(vec![2, 4]);
1976        let values1 = Int32Array::from(vec![10, 20]);
1977        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
1978        let run_ends2 = Int32Array::from(vec![1, 4]);
1979        let values2 = Int32Array::from(vec![30, 40]);
1980        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
1981        let run_ends3 = Int32Array::from(vec![1, 4]);
1982        let values3 = Int32Array::from(vec![50, 60]);
1983        let array3 = RunArray::try_new(&run_ends3, &values3).unwrap();
1984
1985        // Concatenate the arrays
1986        let result = concat(&[&array1, &array2, &array3]).unwrap();
1987        let result_run_array: &arrow_array::RunArray<Int32Type> = result.as_run();
1988
1989        // Check that the result has the correct length
1990        assert_eq!(result_run_array.len(), 12); // 4 + 4 + 4
1991
1992        // Check the run ends
1993        let run_ends = result_run_array.run_ends().values();
1994        assert_eq!(run_ends.len(), 6);
1995        assert_eq!(&[2, 4, 5, 8, 9, 12], run_ends);
1996
1997        // Check the values
1998        let values = result_run_array
1999            .values()
2000            .as_any()
2001            .downcast_ref::<Int32Array>()
2002            .unwrap();
2003        assert_eq!(values.len(), 6);
2004        assert_eq!(&[10, 20, 30, 40, 50, 60], values.values());
2005    }
2006
2007    #[test]
2008    fn test_concat_run_array_with_truncated_run() {
2009        // Create a run array with run ends [2, 5] and values [10, 20]
2010        // Logical: [10, 10, 20, 20, 20]
2011        let run_ends1 = Int32Array::from(vec![2, 5]);
2012        let values1 = Int32Array::from(vec![10, 20]);
2013        let array1 = RunArray::try_new(&run_ends1, &values1).unwrap();
2014        let array1_sliced = array1.slice(0, 3);
2015
2016        let run_ends2 = Int32Array::from(vec![2]);
2017        let values2 = Int32Array::from(vec![30]);
2018        let array2 = RunArray::try_new(&run_ends2, &values2).unwrap();
2019
2020        let result = concat(&[&array1_sliced, &array2]).unwrap();
2021        let result_run_array = result.as_run::<Int32Type>();
2022
2023        // Result should be [10, 10, 20, 30, 30]
2024        // Run ends should be [2, 3, 5]
2025        assert_eq!(result_run_array.len(), 5);
2026        let run_ends = result_run_array.run_ends().values();
2027        let values = result_run_array.values().as_primitive::<Int32Type>();
2028        assert_eq!(values.values(), &[10, 20, 30]);
2029        assert_eq!(&[2, 3, 5], run_ends);
2030    }
2031
2032    /// A single row of a {String -> Int32} map: `None` for a null row, otherwise
2033    /// the list of (key, optional value) entries.
2034    type StringIntMapRow<'a> = Option<Vec<(&'a str, Option<i32>)>>;
2035
2036    /// Helper to build a MapArray of {String -> Int32} from a list of entries per row.
2037    fn build_string_int_map(rows: Vec<StringIntMapRow>) -> MapArray {
2038        let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32ArrayBuilder::new());
2039        for row in rows {
2040            match row {
2041                Some(entries) => {
2042                    for (k, v) in entries {
2043                        builder.keys().append_value(k);
2044                        builder.values().append_option(v);
2045                    }
2046                    builder.append(true).unwrap();
2047                }
2048                None => {
2049                    builder.append(false).unwrap();
2050                }
2051            }
2052        }
2053        builder.finish()
2054    }
2055
2056    #[test]
2057    fn test_concat_map_arrays() {
2058        let map1 = build_string_int_map(vec![
2059            Some(vec![("a", Some(1)), ("b", Some(2))]),
2060            Some(vec![("c", Some(3))]),
2061        ]);
2062        let map2 = build_string_int_map(vec![
2063            Some(vec![("d", Some(4)), ("e", Some(5))]),
2064            None,
2065            Some(vec![("f", Some(6))]),
2066        ]);
2067
2068        let result = concat(&[&map1, &map2]).unwrap();
2069        let result_map = result.as_map();
2070
2071        assert_eq!(result_map.len(), 5);
2072        assert_eq!(result_map.null_count(), 1);
2073
2074        // Check offsets
2075        assert_eq!(result_map.value_offsets(), &[0, 2, 3, 5, 5, 6]);
2076
2077        // Check keys
2078        let keys = result_map.keys().as_string::<i32>();
2079        let expected_keys: Vec<&str> = vec!["a", "b", "c", "d", "e", "f"];
2080        let actual_keys: Vec<&str> = keys.iter().map(|v| v.unwrap()).collect();
2081        assert_eq!(actual_keys, expected_keys);
2082
2083        // Check values
2084        let values = result_map.values().as_primitive::<Int32Type>();
2085        assert_eq!(values.values(), &[1, 2, 3, 4, 5, 6]);
2086    }
2087
2088    #[test]
2089    fn test_concat_map_arrays_sliced() {
2090        let map = build_string_int_map(vec![
2091            Some(vec![("a", Some(1))]),
2092            Some(vec![("b", Some(2)), ("c", Some(3))]),
2093            Some(vec![("d", Some(4))]),
2094            Some(vec![("e", Some(5))]),
2095        ]);
2096
2097        // Slice to get the middle two rows: [("b",2),("c",3)] and [("d",4)]
2098        let sliced = map.slice(1, 2);
2099
2100        let map2 = build_string_int_map(vec![Some(vec![("f", Some(6))])]);
2101
2102        let result = concat(&[&sliced, &map2]).unwrap();
2103        let result_map = result.as_map();
2104
2105        assert_eq!(result_map.len(), 3);
2106        assert_eq!(result_map.value_offsets(), &[0, 2, 3, 4]);
2107
2108        let keys = result_map.keys().as_string::<i32>();
2109        let actual_keys: Vec<&str> = keys.iter().map(|v| v.unwrap()).collect();
2110        assert_eq!(actual_keys, vec!["b", "c", "d", "f"]);
2111    }
2112
2113    #[test]
2114    fn test_concat_map_arrays_with_nulls() {
2115        let map1 = build_string_int_map(vec![Some(vec![("a", Some(1))]), None]);
2116        let map2 = build_string_int_map(vec![None, Some(vec![("b", Some(2))])]);
2117
2118        let result = concat(&[&map1, &map2]).unwrap();
2119        let result_map = result.as_map();
2120
2121        assert_eq!(result_map.len(), 4);
2122        assert_eq!(result_map.null_count(), 2);
2123        assert!(result_map.is_valid(0));
2124        assert!(result_map.is_null(1));
2125        assert!(result_map.is_null(2));
2126        assert!(result_map.is_valid(3));
2127    }
2128
2129    #[test]
2130    fn test_concat_map_arrays_empty_maps() {
2131        let map1 = build_string_int_map(vec![Some(vec![]), Some(vec![("a", Some(1))])]);
2132        let map2 = build_string_int_map(vec![
2133            Some(vec![]),
2134            Some(vec![("b", Some(2)), ("c", Some(3))]),
2135        ]);
2136
2137        let result = concat(&[&map1, &map2]).unwrap();
2138        let result_map = result.as_map();
2139
2140        assert_eq!(result_map.len(), 4);
2141        assert_eq!(result_map.null_count(), 0);
2142        assert_eq!(result_map.value_offsets(), &[0, 0, 1, 1, 3]);
2143    }
2144}