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