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