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