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