1use crate::concat::concat;
21use crate::dictionary::{merge_dictionary_values, should_merge_dictionary_values};
22use arrow_array::ListLikeArray;
23use arrow_array::builder::{BooleanBufferBuilder, PrimitiveBuilder};
24use arrow_array::cast::AsArray;
25use arrow_array::types::*;
26use arrow_array::*;
27use arrow_buffer::bit_mask::set_bits;
28use arrow_buffer::bit_util;
29use arrow_buffer::{
30 ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer, OffsetBuffer,
31};
32use arrow_data::ByteView;
33use arrow_data::transform::MutableArrayData;
34use arrow_schema::{ArrowError, DataType, FieldRef, Fields};
35use std::collections::BTreeMap;
36use std::sync::Arc;
37
38macro_rules! primitive_helper {
39 ($t:ty, $values:ident, $indices:ident, $data_type:ident) => {
40 interleave_primitive::<$t>($values, $indices, $data_type)
41 };
42}
43
44macro_rules! dict_helper {
45 ($t:ty, $values:expr, $indices:expr) => {
46 interleave_dictionaries::<$t>($values, $indices)
47 };
48}
49
50pub fn interleave(
78 values: &[&dyn Array],
79 indices: &[(usize, usize)],
80) -> Result<ArrayRef, ArrowError> {
81 if values.is_empty() {
82 return Err(ArrowError::InvalidArgumentError(
83 "interleave requires input of at least one array".to_string(),
84 ));
85 }
86 let data_type = values[0].data_type();
87
88 for array in values.iter().skip(1) {
89 if array.data_type() != data_type {
90 return Err(ArrowError::InvalidArgumentError(format!(
91 "It is not possible to interleave arrays of different data types ({} and {})",
92 data_type,
93 array.data_type()
94 )));
95 }
96 }
97
98 if indices.is_empty() {
99 return Ok(new_empty_array(data_type));
100 }
101
102 downcast_primitive! {
103 data_type => (primitive_helper, values, indices, data_type),
104 DataType::Boolean => interleave_boolean(values, indices),
105 DataType::Utf8 => interleave_bytes::<Utf8Type>(values, indices),
106 DataType::LargeUtf8 => interleave_bytes::<LargeUtf8Type>(values, indices),
107 DataType::Binary => interleave_bytes::<BinaryType>(values, indices),
108 DataType::LargeBinary => interleave_bytes::<LargeBinaryType>(values, indices),
109 DataType::BinaryView => interleave_views::<BinaryViewType>(values, indices),
110 DataType::Utf8View => interleave_views::<StringViewType>(values, indices),
111 DataType::Dictionary(k, _) => downcast_integer! {
112 k.as_ref() => (dict_helper, values, indices),
113 _ => unreachable!("illegal dictionary key type {k}")
114 },
115 DataType::Struct(fields) => interleave_struct(fields, values, indices),
116 DataType::List(field) => interleave_list::<i32>(values, indices, field),
117 DataType::LargeList(field) => interleave_list::<i64>(values, indices, field),
118 DataType::FixedSizeList(field, size) => interleave_fixed_size_list(values, indices, field, *size),
119 DataType::Map(field, ordered) => interleave_map(values, indices, field, *ordered),
120 DataType::RunEndEncoded(r, _) => match r.data_type() {
121 DataType::Int16 => interleave_run_end::<Int16Type>(values, indices),
122 DataType::Int32 => interleave_run_end::<Int32Type>(values, indices),
123 DataType::Int64 => interleave_run_end::<Int64Type>(values, indices),
124 t => unreachable!("illegal run-end type {t}"),
125 },
126 DataType::ListView(field) => interleave_list_view::<i32>(values, indices, field),
127 DataType::LargeListView(field) => interleave_list_view::<i64>(values, indices, field),
128 _ => interleave_fallback(values, indices)
129 }
130}
131
132struct Interleave<'a, T> {
136 arrays: Vec<&'a T>,
138 nulls: Option<NullBuffer>,
140}
141
142impl<'a, T: Array + 'static> Interleave<'a, T> {
143 fn new(values: &[&'a dyn Array], indices: &'a [(usize, usize)]) -> Self {
144 let mut has_nulls = false;
145 let arrays: Vec<&T> = values
146 .iter()
147 .map(|x| {
148 has_nulls = has_nulls || x.null_count() != 0;
149 x.as_any().downcast_ref().unwrap()
150 })
151 .collect();
152
153 let nulls = match has_nulls {
154 true => {
155 let nulls = BooleanBuffer::collect_bool(indices.len(), |i| {
156 let (a, b) = indices[i];
157 arrays[a].is_valid(b)
158 });
159 Some(nulls.into())
160 }
161 false => None,
162 };
163
164 Self { arrays, nulls }
165 }
166}
167
168fn interleave_primitive<T: ArrowPrimitiveType>(
169 values: &[&dyn Array],
170 indices: &[(usize, usize)],
171 data_type: &DataType,
172) -> Result<ArrayRef, ArrowError> {
173 let interleaved = Interleave::<'_, PrimitiveArray<T>>::new(values, indices);
174 let arrays = &interleaved.arrays;
175 let len = indices.len();
176
177 let mut output = Vec::with_capacity(len);
178 let dst: *mut T::Native = output.as_mut_ptr();
179 let mut base = 0;
180
181 let (chunks, remainder) = indices.as_chunks::<8>();
184 for chunk in chunks {
185 let v0 = arrays[chunk[0].0].value(chunk[0].1);
186 let v1 = arrays[chunk[1].0].value(chunk[1].1);
187 let v2 = arrays[chunk[2].0].value(chunk[2].1);
188 let v3 = arrays[chunk[3].0].value(chunk[3].1);
189 let v4 = arrays[chunk[4].0].value(chunk[4].1);
190 let v5 = arrays[chunk[5].0].value(chunk[5].1);
191 let v6 = arrays[chunk[6].0].value(chunk[6].1);
192 let v7 = arrays[chunk[7].0].value(chunk[7].1);
193
194 debug_assert!(base + 7 < len);
196 unsafe {
197 dst.add(base).write(v0);
198 dst.add(base + 1).write(v1);
199 dst.add(base + 2).write(v2);
200 dst.add(base + 3).write(v3);
201 dst.add(base + 4).write(v4);
202 dst.add(base + 5).write(v5);
203 dst.add(base + 6).write(v6);
204 dst.add(base + 7).write(v7);
205 }
206 base += 8;
207 }
208
209 for idx in remainder {
210 debug_assert!(base < len);
212 unsafe { dst.add(base).write(arrays[idx.0].value(idx.1)) };
213 base += 1;
214 }
215
216 debug_assert_eq!(base, len);
218 unsafe { output.set_len(len) };
219
220 let array = PrimitiveArray::<T>::try_new(output.into(), interleaved.nulls)?;
221 Ok(Arc::new(array.with_data_type(data_type.clone())))
222}
223
224fn interleave_boolean(
225 values: &[&dyn Array],
226 indices: &[(usize, usize)],
227) -> Result<ArrayRef, ArrowError> {
228 let interleaved = Interleave::<'_, BooleanArray>::new(values, indices);
229 let arrays = &interleaved.arrays;
230
231 let mut values_builder = BooleanBufferBuilder::new(indices.len());
232
233 for &(array_idx, row_idx) in indices {
234 values_builder.append(arrays[array_idx].value(row_idx));
235 }
236
237 let array = BooleanArray::new(values_builder.build(), interleaved.nulls);
238 Ok(Arc::new(array))
239}
240
241fn interleave_bytes<T: ByteArrayType>(
242 values: &[&dyn Array],
243 indices: &[(usize, usize)],
244) -> Result<ArrayRef, ArrowError> {
245 let interleaved = Interleave::<'_, GenericByteArray<T>>::new(values, indices);
246
247 let mut capacity = 0;
248 let mut offsets = Vec::with_capacity(indices.len() + 1);
249 offsets.push(T::Offset::from_usize(0).unwrap());
250 for (a, b) in indices {
251 let o = interleaved.arrays[*a].value_offsets();
252 let element_len = o[*b + 1].as_usize() - o[*b].as_usize();
253 capacity += element_len;
254 offsets.push(
255 T::Offset::from_usize(capacity)
256 .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?,
257 );
258 }
259
260 let mut values = Vec::with_capacity(capacity);
261 for (a, b) in indices {
262 values.extend_from_slice(interleaved.arrays[*a].value(*b).as_ref());
263 }
264
265 let array = unsafe {
267 let offsets = OffsetBuffer::new_unchecked(offsets.into());
268 GenericByteArray::<T>::new_unchecked(offsets, values.into(), interleaved.nulls)
269 };
270 Ok(Arc::new(array))
271}
272
273fn interleave_dictionaries<K: ArrowDictionaryKeyType>(
274 arrays: &[&dyn Array],
275 indices: &[(usize, usize)],
276) -> Result<ArrayRef, ArrowError> {
277 let dictionaries: Vec<_> = arrays.iter().map(|x| x.as_dictionary::<K>()).collect();
278 let (should_merge, has_overflow) =
279 should_merge_dictionary_values::<K>(&dictionaries, indices.len());
280 if !should_merge {
281 return if has_overflow {
282 interleave_fallback(arrays, indices)
283 } else {
284 interleave_fallback_dictionary::<K>(&dictionaries, indices)
285 };
286 }
287
288 let masks: Vec<_> = dictionaries
289 .iter()
290 .enumerate()
291 .map(|(a_idx, dictionary)| {
292 let mut key_mask = BooleanBufferBuilder::new_from_buffer(
293 MutableBuffer::new_null(dictionary.len()),
294 dictionary.len(),
295 );
296
297 for (_, key_idx) in indices.iter().filter(|(a, _)| *a == a_idx) {
298 key_mask.set_bit(*key_idx, true);
299 }
300 key_mask.finish()
301 })
302 .collect();
303
304 let merged = merge_dictionary_values(&dictionaries, Some(&masks))?;
305
306 let mut keys = PrimitiveBuilder::<K>::with_capacity(indices.len());
308 for (a, b) in indices {
309 let old_keys: &PrimitiveArray<K> = dictionaries[*a].keys();
310 match old_keys.is_valid(*b) {
311 true => {
312 let old_key = old_keys.values()[*b];
313 keys.append_value(merged.key_mappings[*a][old_key.as_usize()])
314 }
315 false => keys.append_null(),
316 }
317 }
318 let array = unsafe { DictionaryArray::new_unchecked(keys.finish(), merged.values) };
319 Ok(Arc::new(array))
320}
321
322fn interleave_views<T: ByteViewType>(
323 values: &[&dyn Array],
324 indices: &[(usize, usize)],
325) -> Result<ArrayRef, ArrowError> {
326 let interleaved = Interleave::<'_, GenericByteViewArray<T>>::new(values, indices);
327 let mut buffers = Vec::new();
328
329 let mut offsets = Vec::with_capacity(interleaved.arrays.len() + 1);
331 offsets.push(0);
332 let mut total_buffers = 0;
333 for a in &interleaved.arrays {
334 total_buffers += a.data_buffers().len();
335 offsets.push(total_buffers);
336 }
337
338 const UNASSIGNED: u32 = u32::MAX;
343
344 let mut buffer_to_new_index = vec![UNASSIGNED; total_buffers];
346
347 let mut seen_buffers = BTreeMap::new();
349
350 let views: Vec<u128> = indices
351 .iter()
352 .map(|(array_idx, value_idx)| {
353 let array = interleaved.arrays[*array_idx];
354 let view = array.views().get(*value_idx).unwrap();
355 let view_len = *view as u32;
356 if view_len <= 12 {
357 return *view;
358 }
359 let view = ByteView::from(*view);
361 let buffer_to_new_idx = offsets[*array_idx] + view.buffer_index as usize;
362 let new_buffer_idx = &mut buffer_to_new_index[buffer_to_new_idx];
363 if *new_buffer_idx == UNASSIGNED {
364 *new_buffer_idx = push_buffer(
365 &mut seen_buffers,
366 &mut buffers,
367 &array.data_buffers()[view.buffer_index as usize],
368 );
369 }
370 view.with_buffer_index(*new_buffer_idx).as_u128()
371 })
372 .collect();
373
374 let array = unsafe {
375 GenericByteViewArray::<T>::new_unchecked(views.into(), buffers.into(), interleaved.nulls)
376 };
377 Ok(Arc::new(array))
378}
379
380#[inline(never)]
385fn push_buffer(
386 seen: &mut BTreeMap<(*const u8, usize), u32>,
387 buffers: &mut Vec<Buffer>,
388 buffer: &Buffer,
389) -> u32 {
390 *seen
391 .entry((buffer.as_ptr(), buffer.len()))
392 .or_insert_with(|| {
393 buffers.push(buffer.clone());
394 (buffers.len() - 1) as u32
395 })
396}
397
398fn interleave_struct(
399 fields: &Fields,
400 values: &[&dyn Array],
401 indices: &[(usize, usize)],
402) -> Result<ArrayRef, ArrowError> {
403 let interleaved = Interleave::<'_, StructArray>::new(values, indices);
404
405 if fields.is_empty() {
406 let array = StructArray::try_new_with_length(
407 fields.clone(),
408 vec![],
409 interleaved.nulls,
410 indices.len(),
411 )?;
412 return Ok(Arc::new(array));
413 }
414
415 let struct_fields_array: Result<Vec<_>, _> = (0..fields.len())
416 .map(|i| {
417 let field_values: Vec<&dyn Array> = interleaved
418 .arrays
419 .iter()
420 .map(|x| x.column(i).as_ref())
421 .collect();
422 interleave(&field_values, indices)
423 })
424 .collect();
425
426 let struct_array =
427 StructArray::try_new(fields.clone(), struct_fields_array?, interleaved.nulls)?;
428 Ok(Arc::new(struct_array))
429}
430
431fn interleave_list_like_primitive_child<L: ListLikeArray, T: ArrowPrimitiveType>(
432 interleaved: &Interleave<'_, L>,
433 indices: &[(usize, usize)],
434 capacity: usize,
435 data_type: &DataType,
436) -> ArrayRef {
437 let child_arrays: Vec<&PrimitiveArray<T>> = interleaved
438 .arrays
439 .iter()
440 .map(|list| list.values().as_primitive::<T>())
441 .collect();
442
443 let has_child_nulls = child_arrays.iter().any(|a| a.null_count() > 0);
444
445 let mut values: Vec<T::Native> = Vec::with_capacity(capacity);
447 for &(array, row) in indices {
448 let range = interleaved.arrays[array].element_range(row);
449 if !range.is_empty() {
450 values.extend_from_slice(&child_arrays[array].values()[range]);
451 }
452 }
453
454 let nulls = if has_child_nulls {
458 let null_byte_len = bit_util::ceil(capacity, 8);
459 let mut output_null_buf = MutableBuffer::from_len_zeroed(null_byte_len);
460
461 let mut offset_write = 0;
462 let mut output_null_count = 0usize;
463 for &(array, row) in indices {
464 let range = interleaved.arrays[array].element_range(row);
465 let len = range.len();
466 if len > 0 {
467 match child_arrays[array].nulls() {
468 Some(null_buffer) => {
469 output_null_count += set_bits(
470 output_null_buf.as_slice_mut(),
471 null_buffer.validity(),
472 offset_write,
473 null_buffer.offset() + range.start,
474 len,
475 );
476 }
477 None => {
478 let buf = output_null_buf.as_slice_mut();
480 (offset_write..offset_write + len).for_each(|i| bit_util::set_bit(buf, i));
481 }
482 }
483 }
484 offset_write += len;
485 }
486
487 if output_null_count > 0 {
488 let bool_buf = BooleanBuffer::new(output_null_buf.into(), 0, capacity);
489 Some(unsafe { NullBuffer::new_unchecked(bool_buf, output_null_count) })
491 } else {
492 None
493 }
494 } else {
495 None
496 };
497
498 Arc::new(PrimitiveArray::<T>::new(values.into(), nulls).with_data_type(data_type.clone()))
499}
500
501fn interleave_list_like_child<L: ListLikeArray>(
503 interleaved: &Interleave<'_, L>,
504 indices: &[(usize, usize)],
505 capacity: usize,
506) -> Result<ArrayRef, ArrowError> {
507 let mut child_indices = Vec::with_capacity(capacity);
508 for &(array, row) in indices {
509 let range = interleaved.arrays[array].element_range(row);
510 child_indices.extend(range.map(|i| (array, i)));
511 }
512
513 let child_arrays: Vec<&dyn Array> = interleaved
514 .arrays
515 .iter()
516 .map(|list| list.values().as_ref())
517 .collect();
518 interleave(&child_arrays, &child_indices)
519}
520
521fn interleave_list<O: OffsetSizeTrait>(
522 values: &[&dyn Array],
523 indices: &[(usize, usize)],
524 field: &FieldRef,
525) -> Result<ArrayRef, ArrowError> {
526 let interleaved = Interleave::<'_, GenericListArray<O>>::new(values, indices);
527
528 let mut capacity = 0usize;
530 let mut offsets = Vec::with_capacity(indices.len() + 1);
531 offsets.push(O::from_usize(0).unwrap());
532 for (array, row) in indices {
533 let o = interleaved.arrays[*array].value_offsets();
534 let element_len = o[*row + 1].as_usize() - o[*row].as_usize();
535 capacity += element_len;
536 offsets.push(
537 O::from_usize(capacity).ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?,
538 );
539 }
540
541 macro_rules! list_primitive_helper {
543 ($t:ty) => {
544 interleave_list_like_primitive_child::<GenericListArray<O>, $t>(
545 &interleaved,
546 indices,
547 capacity,
548 field.data_type(),
549 )
550 };
551 }
552
553 let child_values = downcast_primitive! {
554 field.data_type() => (list_primitive_helper),
558 _ => {
559 interleave_list_like_child(&interleaved, indices, capacity)?
560 }
561 };
562
563 let offsets = OffsetBuffer::new(offsets.into());
564 let list_array =
565 GenericListArray::<O>::new(field.clone(), offsets, child_values, interleaved.nulls);
566
567 Ok(Arc::new(list_array))
568}
569
570fn interleave_fixed_size_list(
571 values: &[&dyn Array],
572 indices: &[(usize, usize)],
573 field: &FieldRef,
574 size: i32,
575) -> Result<ArrayRef, ArrowError> {
576 let interleaved = Interleave::<'_, FixedSizeListArray>::new(values, indices);
577 let capacity = indices.len() * size as usize;
578
579 macro_rules! fsl_primitive_helper {
580 ($t:ty) => {
581 interleave_list_like_primitive_child::<FixedSizeListArray, $t>(
582 &interleaved,
583 indices,
584 capacity,
585 field.data_type(),
586 )
587 };
588 }
589
590 let interleaved_values = downcast_primitive! {
591 field.data_type() => (fsl_primitive_helper),
592 _ => {
593 interleave_list_like_child(&interleaved, indices, capacity)?
594 }
595 };
596
597 let array = FixedSizeListArray::try_new_with_length(
598 field.clone(),
599 size,
600 interleaved_values,
601 interleaved.nulls,
602 indices.len(),
603 )?;
604 Ok(Arc::new(array))
605}
606
607fn interleave_map(
608 values: &[&dyn Array],
609 indices: &[(usize, usize)],
610 field: &FieldRef,
611 ordered: bool,
612) -> Result<ArrayRef, ArrowError> {
613 let interleaved = Interleave::<'_, MapArray>::new(values, indices);
614
615 let mut capacity = 0usize;
616 let mut offsets = Vec::with_capacity(indices.len() + 1);
617 offsets.push(0i32);
618 for &(array, row) in indices {
619 let o = interleaved.arrays[array].value_offsets();
620 let element_len = (o[row + 1] - o[row]) as usize;
621 capacity += element_len;
622 offsets
623 .push(i32::try_from(capacity).map_err(|_| ArrowError::OffsetOverflowError(capacity))?);
624 }
625
626 let mut child_indices = Vec::with_capacity(capacity);
627 for &(array, row) in indices {
628 let o = interleaved.arrays[array].value_offsets();
629 let start = o[row] as usize;
630 let end = o[row + 1] as usize;
631 child_indices.extend((start..end).map(|i| (array, i)));
632 }
633
634 let entries_arrays: Vec<&dyn Array> = interleaved
635 .arrays
636 .iter()
637 .map(|m| m.entries() as &dyn Array)
638 .collect();
639 let interleaved_entries = interleave(&entries_arrays, &child_indices)?;
640
641 let offsets = OffsetBuffer::new(offsets.into());
642 let entries = interleaved_entries.as_struct().clone();
643 let array = MapArray::new(field.clone(), offsets, entries, interleaved.nulls, ordered);
644 Ok(Arc::new(array))
645}
646
647fn interleave_run_end<R: RunEndIndexType>(
649 values: &[&dyn Array],
650 indices: &[(usize, usize)],
651) -> Result<ArrayRef, ArrowError> {
652 if indices.is_empty() {
653 return Ok(new_empty_array(values[0].data_type()));
654 }
655
656 let n = indices.len();
657 R::Native::from_usize(n).ok_or_else(|| {
658 ArrowError::ComputeError(format!(
659 "interleave_run_end: output length {n} does not fit run-end type"
660 ))
661 })?;
662
663 let runs: Vec<&RunArray<R>> = values.iter().map(|a| a.as_run::<R>()).collect();
664 let value_arrays: Vec<&dyn Array> = runs.iter().map(|r| r.values().as_ref()).collect();
665
666 let mut phys_pairs: Vec<(usize, usize)> = vec![(0, 0); n];
669 let mut grouped: Vec<(Vec<R::Native>, Vec<usize>)> =
670 (0..runs.len()).map(|_| (Vec::new(), Vec::new())).collect();
671 for (out_pos, &(arr, row)) in indices.iter().enumerate() {
672 let row = R::Native::from_usize(row).ok_or_else(|| {
673 ArrowError::InvalidArgumentError(format!(
674 "interleave_run_end: row index {row} not representable as run-end type {}",
675 R::DATA_TYPE
676 ))
677 })?;
678 grouped[arr].0.push(row);
679 grouped[arr].1.push(out_pos);
680 }
681 for (arr_idx, (logical_rows, out_positions)) in grouped.into_iter().enumerate() {
682 let phys = runs[arr_idx].get_physical_indices(&logical_rows)?;
683 for (p, out_pos) in phys.iter().zip(out_positions.iter()) {
684 phys_pairs[*out_pos] = (arr_idx, *p);
685 }
686 }
687
688 let mut run_ends_buf: Vec<R::Native> = Vec::with_capacity(n);
693 let mut dedup_pairs: Vec<(usize, usize)> = Vec::with_capacity(n);
694 dedup_pairs.push(phys_pairs[0]);
695 for i in 1..n {
696 if phys_pairs[i] != phys_pairs[i - 1] {
697 run_ends_buf.push(R::Native::from_usize(i).unwrap());
698 dedup_pairs.push(phys_pairs[i]);
699 }
700 }
701 run_ends_buf.push(R::Native::from_usize(n).unwrap());
702
703 let taken_values = interleave(&value_arrays, &dedup_pairs)?;
704 let run_ends = PrimitiveArray::<R>::from_iter_values(run_ends_buf);
705
706 Ok(Arc::new(RunArray::<R>::try_new(
707 &run_ends,
708 taken_values.as_ref(),
709 )?))
710}
711
712fn interleave_list_view<O: OffsetSizeTrait>(
713 values: &[&dyn Array],
714 indices: &[(usize, usize)],
715 field: &FieldRef,
716) -> Result<ArrayRef, ArrowError> {
717 let interleaved = Interleave::<'_, GenericListViewArray<O>>::new(values, indices);
718
719 let concat_cost: usize = interleaved.arrays.iter().map(|lv| lv.values().len()).sum();
724 let per_row_cost: usize = indices
725 .iter()
726 .map(|&(a, r)| interleaved.arrays[a].sizes()[r].as_usize())
727 .sum();
728
729 if per_row_cost <= concat_cost {
730 interleave_list_view_copy::<O>(&interleaved, indices, field)
731 } else {
732 interleave_list_view_concat::<O>(&interleaved, indices, field)
733 }
734}
735
736fn interleave_list_view_copy<O: OffsetSizeTrait>(
738 interleaved: &Interleave<'_, GenericListViewArray<O>>,
739 indices: &[(usize, usize)],
740 field: &FieldRef,
741) -> Result<ArrayRef, ArrowError> {
742 let mut capacity = 0usize;
743 let mut offsets = Vec::with_capacity(indices.len());
744 let mut sizes = Vec::with_capacity(indices.len());
745 for &(array_idx, row_idx) in indices {
746 let list = interleaved.arrays[array_idx];
747 let size = list.sizes()[row_idx].as_usize();
748 offsets.push(
749 O::from_usize(capacity).ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?,
750 );
751 sizes.push(O::from_usize(size).ok_or_else(|| ArrowError::OffsetOverflowError(size))?);
752 capacity += size;
753 }
754
755 let child_data: Vec<_> = interleaved
756 .arrays
757 .iter()
758 .map(|list| list.values().to_data())
759 .collect();
760 let child_data_refs: Vec<_> = child_data.iter().collect();
761 let mut mutable_child = MutableArrayData::new(child_data_refs, false, capacity);
762 for &(array_idx, row_idx) in indices {
763 let list = interleaved.arrays[array_idx];
764 let start = list.offsets()[row_idx].as_usize();
765 let size = list.sizes()[row_idx].as_usize();
766 if size > 0 {
767 mutable_child.try_extend(array_idx, start, start + size)?;
768 }
769 }
770
771 Ok(Arc::new(GenericListViewArray::<O>::new(
772 field.clone(),
773 offsets.into(),
774 sizes.into(),
775 make_array(mutable_child.freeze()),
776 interleaved.nulls.clone(),
777 )))
778}
779
780fn interleave_list_view_concat<O: OffsetSizeTrait>(
783 interleaved: &Interleave<'_, GenericListViewArray<O>>,
784 indices: &[(usize, usize)],
785 field: &FieldRef,
786) -> Result<ArrayRef, ArrowError> {
787 let child_arrays: Vec<&dyn Array> = interleaved
788 .arrays
789 .iter()
790 .map(|lv| lv.values().as_ref())
791 .collect();
792 let mut base_offsets = Vec::with_capacity(interleaved.arrays.len());
793 let mut running = 0usize;
794 for lv in &interleaved.arrays {
795 base_offsets.push(running);
796 running += lv.values().len();
797 }
798 let combined_values = concat(&child_arrays)?;
799
800 let mut new_offsets = Vec::with_capacity(indices.len());
801 let mut new_sizes = Vec::with_capacity(indices.len());
802 for &(array_idx, row_idx) in indices {
803 let lv = interleaved.arrays[array_idx];
804 let adjusted = lv.offsets()[row_idx].as_usize() + base_offsets[array_idx];
805 new_offsets.push(
806 O::from_usize(adjusted).ok_or_else(|| ArrowError::OffsetOverflowError(adjusted))?,
807 );
808 new_sizes.push(lv.sizes()[row_idx]);
809 }
810
811 Ok(Arc::new(GenericListViewArray::<O>::new(
812 field.clone(),
813 new_offsets.into(),
814 new_sizes.into(),
815 combined_values,
816 interleaved.nulls.clone(),
817 )))
818}
819
820fn interleave_fallback(
822 values: &[&dyn Array],
823 indices: &[(usize, usize)],
824) -> Result<ArrayRef, ArrowError> {
825 let arrays: Vec<_> = values.iter().map(|x| x.to_data()).collect();
826 let arrays: Vec<_> = arrays.iter().collect();
827 let mut array_data = MutableArrayData::try_new(arrays, false, indices.len())?;
828
829 let mut cur_array = indices[0].0;
830 let mut start_row_idx = indices[0].1;
831 let mut end_row_idx = start_row_idx + 1;
832
833 for (array, row) in indices.iter().skip(1).copied() {
834 if array == cur_array && row == end_row_idx {
835 end_row_idx += 1;
837 continue;
838 }
839
840 array_data.try_extend(cur_array, start_row_idx, end_row_idx)?;
842
843 cur_array = array;
845 start_row_idx = row;
846 end_row_idx = start_row_idx + 1;
847 }
848
849 array_data.try_extend(cur_array, start_row_idx, end_row_idx)?;
851 Ok(make_array(array_data.freeze()))
852}
853
854fn interleave_fallback_dictionary<K: ArrowDictionaryKeyType>(
864 dictionaries: &[&DictionaryArray<K>],
865 indices: &[(usize, usize)],
866) -> Result<ArrayRef, ArrowError> {
867 let relative_offsets: Vec<usize> = dictionaries
868 .iter()
869 .scan(0usize, |offset, dict| {
870 let current = *offset;
871 *offset += dict.values().len();
872 Some(current)
873 })
874 .collect();
875 let all_values: Vec<&dyn Array> = dictionaries.iter().map(|d| d.values().as_ref()).collect();
876 let concatenated_values = concat(&all_values)?;
877
878 let any_nulls = dictionaries.iter().any(|d| d.keys().nulls().is_some());
879 let (new_keys, nulls) = if any_nulls {
880 let mut has_nulls = false;
881 let new_keys: Vec<K::Native> = indices
882 .iter()
883 .map(|(array, row)| {
884 let old_keys = dictionaries[*array].keys();
885 if old_keys.is_valid(*row) {
886 let old_key = old_keys.values()[*row].as_usize();
887 K::Native::from_usize(relative_offsets[*array] + old_key)
888 .expect("key overflow should be checked by caller")
889 } else {
890 has_nulls = true;
891 K::Native::ZERO
892 }
893 })
894 .collect();
895
896 let nulls = if has_nulls {
897 let null_buffer = BooleanBuffer::collect_bool(indices.len(), |i| {
898 let (array, row) = indices[i];
899 dictionaries[array].keys().is_valid(row)
900 });
901 Some(NullBuffer::new(null_buffer))
902 } else {
903 None
904 };
905 (new_keys, nulls)
906 } else {
907 let new_keys: Vec<K::Native> = indices
908 .iter()
909 .map(|(array, row)| {
910 let old_key = dictionaries[*array].keys().values()[*row].as_usize();
911 K::Native::from_usize(relative_offsets[*array] + old_key)
912 .expect("key overflow should be checked by caller")
913 })
914 .collect();
915 (new_keys, None)
916 };
917
918 let keys_array = PrimitiveArray::<K>::new(new_keys.into(), nulls);
919 let array = unsafe { DictionaryArray::new_unchecked(keys_array, concatenated_values) };
921 Ok(Arc::new(array))
922}
923
924pub fn interleave_record_batch(
969 record_batches: &[&RecordBatch],
970 indices: &[(usize, usize)],
971) -> Result<RecordBatch, ArrowError> {
972 let schema = record_batches[0].schema();
973 let columns = (0..schema.fields().len())
974 .map(|i| {
975 let column_values: Vec<&dyn Array> = record_batches
976 .iter()
977 .map(|batch| batch.column(i).as_ref())
978 .collect();
979 interleave(&column_values, indices)
980 })
981 .collect::<Result<Vec<_>, _>>()?;
982 RecordBatch::try_new(schema, columns)
983}
984
985#[cfg(test)]
986mod tests {
987 use super::*;
988 use arrow_array::Int32RunArray;
989 use arrow_array::builder::{
990 GenericListBuilder, Int32Builder, PrimitiveBuilder, PrimitiveRunBuilder,
991 };
992 use arrow_array::types::{Decimal128Type, Int8Type, TimestampMicrosecondType};
993 use arrow_buffer::ScalarBuffer;
994 use arrow_schema::{Field, TimeUnit};
995
996 #[test]
997 fn test_primitive() {
998 let a = Int32Array::from_iter_values([1, 2, 3, 4]);
999 let b = Int32Array::from_iter_values([5, 6, 7]);
1000 let c = Int32Array::from_iter_values([8, 9, 10]);
1001 let values = interleave(&[&a, &b, &c], &[(0, 3), (0, 3), (2, 2), (2, 0), (1, 1)]).unwrap();
1002 let v = values.as_primitive::<Int32Type>();
1003 assert_eq!(v.values(), &[4, 4, 10, 8, 6]);
1004 }
1005
1006 #[test]
1007 fn test_primitive_nulls() {
1008 let a = Int32Array::from_iter_values([1, 2, 3, 4]);
1009 let b = Int32Array::from_iter([Some(1), Some(4), None]);
1010 let values = interleave(&[&a, &b], &[(0, 1), (1, 2), (1, 2), (0, 3), (0, 2)]).unwrap();
1011 let v: Vec<_> = values.as_primitive::<Int32Type>().into_iter().collect();
1012 assert_eq!(&v, &[Some(2), None, None, Some(4), Some(3)])
1013 }
1014
1015 #[test]
1016 fn test_interleave_boolean() {
1017 let a = BooleanArray::from(vec![Some(true), Some(false), None, Some(true)]);
1018 let b = BooleanArray::from(vec![Some(false), Some(true), Some(true), None]);
1019
1020 let values = interleave(&[&a, &b], &[(0, 0), (1, 1), (0, 2), (1, 3), (1, 2)]).unwrap();
1021
1022 let result = values.as_boolean();
1023
1024 assert_eq!(
1025 result.iter().collect::<Vec<_>>(),
1026 vec![Some(true), Some(true), None, None, Some(true),]
1027 );
1028 }
1029
1030 #[test]
1031 fn test_interleave_boolean_sliced() {
1032 let original = BooleanArray::from(vec![
1033 Some(false),
1034 Some(true),
1035 Some(false),
1036 Some(true),
1037 None,
1038 Some(false),
1039 Some(true),
1040 ]);
1041
1042 let sliced_array = original.slice(1, 5);
1043 let sliced = sliced_array
1044 .as_any()
1045 .downcast_ref::<BooleanArray>()
1046 .unwrap();
1047 let other = BooleanArray::from(vec![Some(false), Some(false), Some(true)]);
1048
1049 let values =
1050 interleave(&[sliced, &other], &[(0, 0), (1, 2), (0, 3), (0, 4), (1, 0)]).unwrap();
1051
1052 let result = values.as_boolean();
1053
1054 assert_eq!(
1055 result.iter().collect::<Vec<_>>(),
1056 vec![Some(true), Some(true), None, Some(false), Some(false),]
1057 );
1058 }
1059
1060 #[test]
1061 fn test_interleave_boolean_fallback_matches_specialized() {
1062 let a = BooleanArray::from_iter(
1063 (0..1024).map(|i| if i % 10 == 0 { None } else { Some(i % 2 == 0) }),
1064 );
1065 let b = BooleanArray::from_iter((0..1024).map(|i| Some(i % 3 == 0)));
1066
1067 let indices: Vec<_> = (0..1024).map(|i| (i % 2, (i * 37) % 1024)).collect();
1068
1069 let values: [&dyn Array; 2] = [&a, &b];
1070
1071 let specialized = interleave_boolean(&values, &indices).unwrap();
1072 let fallback = interleave_fallback(&values, &indices).unwrap();
1073
1074 assert_eq!(specialized.as_ref(), fallback.as_ref());
1075 }
1076
1077 #[test]
1078 fn test_primitive_empty() {
1079 let a = Int32Array::from_iter_values([1, 2, 3, 4]);
1080 let v = interleave(&[&a], &[]).unwrap();
1081 assert!(v.is_empty());
1082 assert_eq!(v.data_type(), &DataType::Int32);
1083 }
1084
1085 #[test]
1086 fn test_strings() {
1087 let a = StringArray::from_iter_values(["a", "b", "c"]);
1088 let b = StringArray::from_iter_values(["hello", "world", "foo"]);
1089 let values = interleave(&[&a, &b], &[(0, 2), (0, 2), (1, 0), (1, 1), (0, 1)]).unwrap();
1090 let v = values.as_string::<i32>();
1091 let values: Vec<_> = v.into_iter().collect();
1092 assert_eq!(
1093 &values,
1094 &[
1095 Some("c"),
1096 Some("c"),
1097 Some("hello"),
1098 Some("world"),
1099 Some("b")
1100 ]
1101 )
1102 }
1103
1104 #[test]
1105 fn test_interleave_dictionary() {
1106 let a = DictionaryArray::<Int32Type>::from_iter(["a", "b", "c", "a", "b"]);
1107 let b = DictionaryArray::<Int32Type>::from_iter(["a", "c", "a", "c", "a"]);
1108
1109 let values =
1111 interleave(&[&a, &b], &[(0, 2), (0, 2), (0, 2), (1, 0), (1, 1), (0, 1)]).unwrap();
1112 let v = values.as_dictionary::<Int32Type>();
1113 assert_eq!(v.values().len(), 5);
1114
1115 let vc = v.downcast_dict::<StringArray>().unwrap();
1116 let collected: Vec<_> = vc.into_iter().map(Option::unwrap).collect();
1117 assert_eq!(&collected, &["c", "c", "c", "a", "c", "b"]);
1118
1119 let values = interleave(&[&a, &b], &[(0, 2), (0, 2), (1, 1)]).unwrap();
1121 let v = values.as_dictionary::<Int32Type>();
1122 assert_eq!(v.values().len(), 1);
1123
1124 let vc = v.downcast_dict::<StringArray>().unwrap();
1125 let collected: Vec<_> = vc.into_iter().map(Option::unwrap).collect();
1126 assert_eq!(&collected, &["c", "c", "c"]);
1127 }
1128
1129 #[test]
1130 fn test_interleave_dictionary_nulls() {
1131 let input_1_keys = Int32Array::from_iter_values([0, 2, 1, 3]);
1132 let input_1_values = StringArray::from(vec![Some("foo"), None, Some("bar"), Some("fiz")]);
1133 let input_1 = DictionaryArray::new(input_1_keys, Arc::new(input_1_values));
1134 let input_2: DictionaryArray<Int32Type> = vec![None].into_iter().collect();
1135
1136 let expected = vec![Some("fiz"), None, None, Some("foo")];
1137
1138 let values = interleave(
1139 &[&input_1 as _, &input_2 as _],
1140 &[(0, 3), (0, 2), (1, 0), (0, 0)],
1141 )
1142 .unwrap();
1143 let dictionary = values.as_dictionary::<Int32Type>();
1144 let actual: Vec<Option<&str>> = dictionary
1145 .downcast_dict::<StringArray>()
1146 .unwrap()
1147 .into_iter()
1148 .collect();
1149
1150 assert_eq!(actual, expected);
1151 }
1152
1153 #[test]
1154 fn test_interleave_dictionary_overflow_same_values() {
1155 let values: ArrayRef = Arc::new(StringArray::from_iter_values(
1156 (0..50).map(|i| format!("v{i}")),
1157 ));
1158
1159 let dict1 = DictionaryArray::<Int8Type>::new(
1165 Int8Array::from_iter_values([0, 1, 2]),
1166 values.clone(),
1167 );
1168 let dict2 = DictionaryArray::<Int8Type>::new(
1169 Int8Array::from_iter_values([0, 1, 2]),
1170 values.clone(),
1171 );
1172 let dict3 =
1173 DictionaryArray::<Int8Type>::new(Int8Array::from_iter_values([49]), values.clone());
1174
1175 let indices = &[(0, 0), (1, 0), (2, 0)];
1176 let result = interleave(&[&dict1, &dict2, &dict3], indices).unwrap();
1177
1178 let dict_result = result.as_dictionary::<Int8Type>();
1179 let string_result: Vec<_> = dict_result
1180 .downcast_dict::<StringArray>()
1181 .unwrap()
1182 .into_iter()
1183 .map(|x| x.unwrap())
1184 .collect();
1185 assert_eq!(string_result, vec!["v0", "v0", "v49"]);
1186 }
1187
1188 fn test_interleave_lists<O: OffsetSizeTrait>() {
1189 let mut a = GenericListBuilder::<O, _>::new(Int32Builder::new());
1191 a.values().append_value(1);
1192 a.values().append_value(2);
1193 a.append(true);
1194 a.append(false);
1195 a.values().append_value(3);
1196 a.append(true);
1197 let a = a.finish();
1198
1199 let mut b = GenericListBuilder::<O, _>::new(Int32Builder::new());
1201 b.values().append_value(4);
1202 b.append(true);
1203 b.append(false);
1204 b.values().append_value(5);
1205 b.values().append_value(6);
1206 b.values().append_null();
1207 b.append(true);
1208 let b = b.finish();
1209
1210 let values = interleave(&[&a, &b], &[(0, 2), (0, 1), (1, 0), (1, 2), (1, 1)]).unwrap();
1211 let v = values
1212 .as_any()
1213 .downcast_ref::<GenericListArray<O>>()
1214 .unwrap();
1215
1216 let mut expected = GenericListBuilder::<O, _>::new(Int32Builder::new());
1218 expected.values().append_value(3);
1219 expected.append(true);
1220 expected.append(false);
1221 expected.values().append_value(4);
1222 expected.append(true);
1223 expected.values().append_value(5);
1224 expected.values().append_value(6);
1225 expected.values().append_null();
1226 expected.append(true);
1227 expected.append(false);
1228 let expected = expected.finish();
1229
1230 assert_eq!(v, &expected);
1231 }
1232
1233 #[test]
1234 fn test_lists() {
1235 test_interleave_lists::<i32>();
1236 }
1237
1238 #[test]
1239 fn test_large_lists() {
1240 test_interleave_lists::<i64>();
1241 }
1242
1243 type ListRow<T> = Option<Vec<Option<<T as ArrowPrimitiveType>::Native>>>;
1246
1247 fn list_of_primitive<O: OffsetSizeTrait, T: ArrowPrimitiveType>(
1250 data_type: &DataType,
1251 rows: &[ListRow<T>],
1252 ) -> GenericListArray<O> {
1253 let mut builder = GenericListBuilder::<O, _>::new(
1254 PrimitiveBuilder::<T>::new().with_data_type(data_type.clone()),
1255 );
1256 for row in rows {
1257 match row {
1258 Some(items) => {
1259 items
1260 .iter()
1261 .for_each(|v| builder.values().append_option(*v));
1262 builder.append(true);
1263 }
1264 None => builder.append(false),
1265 }
1266 }
1267 builder.finish()
1268 }
1269
1270 fn check_interleave_list_primitive<O: OffsetSizeTrait, T: ArrowPrimitiveType>(
1273 data_type: &DataType,
1274 inputs: &[&[ListRow<T>]],
1275 indices: &[(usize, usize)],
1276 expected: &[ListRow<T>],
1277 ) {
1278 let arrays: Vec<_> = inputs
1279 .iter()
1280 .map(|rows| list_of_primitive::<O, T>(data_type, rows))
1281 .collect();
1282 let refs: Vec<&dyn Array> = arrays.iter().map(|a| a as &dyn Array).collect();
1283
1284 let values = interleave(&refs, indices).unwrap();
1285 let v = values
1286 .as_any()
1287 .downcast_ref::<GenericListArray<O>>()
1288 .unwrap();
1289
1290 assert_eq!(v, &list_of_primitive::<O, T>(data_type, expected));
1291 assert_eq!(v.values().data_type(), data_type);
1294 }
1295
1296 fn test_interleave_lists_decimal<O: OffsetSizeTrait>() {
1297 check_interleave_list_primitive::<O, Decimal128Type>(
1299 &DataType::Decimal128(20, 3),
1300 &[
1301 &[
1302 Some(vec![Some(1), Some(2)]),
1303 None,
1304 Some(vec![Some(3), None]),
1305 ], &[Some(vec![Some(4)]), Some(vec![Some(5), Some(6)])], ],
1308 &[(0, 2), (0, 1), (1, 0), (1, 1)],
1309 &[
1310 Some(vec![Some(3), None]),
1311 None,
1312 Some(vec![Some(4)]),
1313 Some(vec![Some(5), Some(6)]),
1314 ],
1315 );
1316 }
1317
1318 #[test]
1319 fn test_lists_decimal() {
1320 test_interleave_lists_decimal::<i32>();
1321 test_interleave_lists_decimal::<i64>();
1322 }
1323
1324 fn test_interleave_lists_timestamp_tz<O: OffsetSizeTrait>() {
1325 check_interleave_list_primitive::<O, TimestampMicrosecondType>(
1327 &DataType::Timestamp(TimeUnit::Microsecond, Some("+08:00".into())),
1328 &[&[Some(vec![Some(1), Some(2)]), Some(vec![Some(3)])]],
1329 &[(0, 1), (0, 0)],
1330 &[Some(vec![Some(3)]), Some(vec![Some(1), Some(2)])],
1331 );
1332 }
1333
1334 #[test]
1335 fn test_lists_timestamp_tz() {
1336 test_interleave_lists_timestamp_tz::<i32>();
1337 test_interleave_lists_timestamp_tz::<i64>();
1338 }
1339
1340 fn test_interleave_list_views<O: OffsetSizeTrait>() {
1341 let mut a = GenericListBuilder::<O, _>::new(Int32Builder::new());
1343 a.values().append_value(1);
1344 a.values().append_value(2);
1345 a.append(true);
1346 a.append(false);
1347 a.values().append_value(3);
1348 a.append(true);
1349 let a: GenericListViewArray<O> = a.finish().into();
1350
1351 let mut b = GenericListBuilder::<O, _>::new(Int32Builder::new());
1353 b.values().append_value(4);
1354 b.append(true);
1355 b.append(false);
1356 b.values().append_value(5);
1357 b.values().append_value(6);
1358 b.values().append_null();
1359 b.append(true);
1360 let b: GenericListViewArray<O> = b.finish().into();
1361
1362 let values = interleave(&[&a, &b], &[(0, 2), (0, 1), (1, 0), (1, 2), (1, 1)]).unwrap();
1363 let v = values
1364 .as_any()
1365 .downcast_ref::<GenericListViewArray<O>>()
1366 .unwrap();
1367
1368 let mut expected = GenericListBuilder::<O, _>::new(Int32Builder::new());
1370 expected.values().append_value(3);
1371 expected.append(true);
1372 expected.append(false);
1373 expected.values().append_value(4);
1374 expected.append(true);
1375 expected.values().append_value(5);
1376 expected.values().append_value(6);
1377 expected.values().append_null();
1378 expected.append(true);
1379 expected.append(false);
1380 let expected: GenericListViewArray<O> = expected.finish().into();
1381
1382 assert_eq!(v, &expected);
1383 }
1384
1385 #[test]
1386 fn test_list_views() {
1387 test_interleave_list_views::<i32>();
1388 }
1389
1390 #[test]
1391 fn test_large_list_views() {
1392 test_interleave_list_views::<i64>();
1393 }
1394
1395 #[test]
1396 fn test_interleave_list_view_overlapping() {
1397 let field = Arc::new(Field::new_list_field(DataType::Int64, false));
1398
1399 let lv_a = ListViewArray::new(
1403 Arc::clone(&field),
1404 ScalarBuffer::from(vec![0i32, 0, 0, 0, 0, 5, 5, 5, 5, 5]),
1405 ScalarBuffer::from(vec![5i32; 10]),
1406 Arc::new(Int64Array::from_iter_values(0..10)),
1407 None,
1408 );
1409
1410 let lv_b = ListViewArray::new(
1414 Arc::clone(&field),
1415 ScalarBuffer::from(vec![0i32, 0, 0, 0, 3, 3, 3, 3]),
1416 ScalarBuffer::from(vec![3i32; 8]),
1417 Arc::new(Int64Array::from_iter_values(100..106)),
1418 None,
1419 );
1420
1421 let indices: Vec<(usize, usize)> = vec![
1422 (0, 0),
1423 (1, 0),
1424 (0, 5),
1425 (1, 4),
1426 (0, 1),
1427 (1, 1),
1428 (0, 6),
1429 (1, 5),
1430 ];
1431 let result = interleave(&[&lv_a as &dyn Array, &lv_b as &dyn Array], &indices).unwrap();
1432 result
1433 .to_data()
1434 .validate_full()
1435 .expect("result must be valid");
1436
1437 let result_lv = result.as_list_view::<i32>();
1438 assert_eq!(result_lv.len(), 8);
1439 assert_eq!(
1440 result_lv.value(0).as_primitive::<Int64Type>().values(),
1441 &[0, 1, 2, 3, 4]
1442 );
1443 assert_eq!(
1444 result_lv.value(1).as_primitive::<Int64Type>().values(),
1445 &[100, 101, 102]
1446 );
1447 assert_eq!(
1448 result_lv.value(2).as_primitive::<Int64Type>().values(),
1449 &[5, 6, 7, 8, 9]
1450 );
1451 assert_eq!(
1452 result_lv.value(3).as_primitive::<Int64Type>().values(),
1453 &[103, 104, 105]
1454 );
1455
1456 let total_input_elements = lv_a.values().len() + lv_b.values().len();
1459 assert_eq!(result_lv.values().len(), total_input_elements);
1460 }
1461
1462 #[test]
1463 fn test_struct_without_nulls() {
1464 let fields = Fields::from(vec![
1465 Field::new("number_col", DataType::Int32, false),
1466 Field::new("string_col", DataType::Utf8, false),
1467 ]);
1468 let a = {
1469 let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
1470 let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);
1471
1472 StructArray::try_new(
1473 fields.clone(),
1474 vec![Arc::new(number_col), Arc::new(string_col)],
1475 None,
1476 )
1477 .unwrap()
1478 };
1479
1480 let b = {
1481 let number_col = Int32Array::from_iter_values([5, 6, 7]);
1482 let string_col = StringArray::from_iter_values(["hello", "world", "foo"]);
1483
1484 StructArray::try_new(
1485 fields.clone(),
1486 vec![Arc::new(number_col), Arc::new(string_col)],
1487 None,
1488 )
1489 .unwrap()
1490 };
1491
1492 let c = {
1493 let number_col = Int32Array::from_iter_values([8, 9, 10]);
1494 let string_col = StringArray::from_iter_values(["x", "y", "z"]);
1495
1496 StructArray::try_new(
1497 fields.clone(),
1498 vec![Arc::new(number_col), Arc::new(string_col)],
1499 None,
1500 )
1501 .unwrap()
1502 };
1503
1504 let values = interleave(&[&a, &b, &c], &[(0, 3), (0, 3), (2, 2), (2, 0), (1, 1)]).unwrap();
1505 let values_struct = values.as_struct();
1506 assert_eq!(values_struct.data_type(), &DataType::Struct(fields));
1507 assert_eq!(values_struct.null_count(), 0);
1508
1509 let values_number = values_struct.column(0).as_primitive::<Int32Type>();
1510 assert_eq!(values_number.values(), &[4, 4, 10, 8, 6]);
1511 let values_string = values_struct.column(1).as_string::<i32>();
1512 let values_string: Vec<_> = values_string.into_iter().collect();
1513 assert_eq!(
1514 &values_string,
1515 &[Some("d"), Some("d"), Some("z"), Some("x"), Some("world")]
1516 );
1517 }
1518
1519 #[test]
1520 fn test_struct_with_nulls_in_values() {
1521 let fields = Fields::from(vec![
1522 Field::new("number_col", DataType::Int32, true),
1523 Field::new("string_col", DataType::Utf8, true),
1524 ]);
1525 let a = {
1526 let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
1527 let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);
1528
1529 StructArray::try_new(
1530 fields.clone(),
1531 vec![Arc::new(number_col), Arc::new(string_col)],
1532 None,
1533 )
1534 .unwrap()
1535 };
1536
1537 let b = {
1538 let number_col = Int32Array::from_iter([Some(1), Some(4), None]);
1539 let string_col = StringArray::from(vec![Some("hello"), None, Some("foo")]);
1540
1541 StructArray::try_new(
1542 fields.clone(),
1543 vec![Arc::new(number_col), Arc::new(string_col)],
1544 None,
1545 )
1546 .unwrap()
1547 };
1548
1549 let values = interleave(&[&a, &b], &[(0, 1), (1, 2), (1, 2), (0, 3), (1, 1)]).unwrap();
1550 let values_struct = values.as_struct();
1551 assert_eq!(values_struct.data_type(), &DataType::Struct(fields));
1552
1553 assert_eq!(values_struct.null_count(), 0);
1555
1556 let values_number: Vec<_> = values_struct
1557 .column(0)
1558 .as_primitive::<Int32Type>()
1559 .into_iter()
1560 .collect();
1561 assert_eq!(values_number, &[Some(2), None, None, Some(4), Some(4)]);
1562
1563 let values_string = values_struct.column(1).as_string::<i32>();
1564 let values_string: Vec<_> = values_string.into_iter().collect();
1565 assert_eq!(
1566 &values_string,
1567 &[Some("b"), Some("foo"), Some("foo"), Some("d"), None]
1568 );
1569 }
1570
1571 #[test]
1572 fn test_struct_with_nulls() {
1573 let fields = Fields::from(vec![
1574 Field::new("number_col", DataType::Int32, false),
1575 Field::new("string_col", DataType::Utf8, false),
1576 ]);
1577 let a = {
1578 let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
1579 let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);
1580
1581 StructArray::try_new(
1582 fields.clone(),
1583 vec![Arc::new(number_col), Arc::new(string_col)],
1584 None,
1585 )
1586 .unwrap()
1587 };
1588
1589 let b = {
1590 let number_col = Int32Array::from_iter_values([5, 6, 7]);
1591 let string_col = StringArray::from_iter_values(["hello", "world", "foo"]);
1592
1593 StructArray::try_new(
1594 fields.clone(),
1595 vec![Arc::new(number_col), Arc::new(string_col)],
1596 Some(NullBuffer::from(&[true, false, true])),
1597 )
1598 .unwrap()
1599 };
1600
1601 let c = {
1602 let number_col = Int32Array::from_iter_values([8, 9, 10]);
1603 let string_col = StringArray::from_iter_values(["x", "y", "z"]);
1604
1605 StructArray::try_new(
1606 fields.clone(),
1607 vec![Arc::new(number_col), Arc::new(string_col)],
1608 None,
1609 )
1610 .unwrap()
1611 };
1612
1613 let values = interleave(&[&a, &b, &c], &[(0, 3), (0, 3), (2, 2), (1, 1), (2, 0)]).unwrap();
1614 let values_struct = values.as_struct();
1615 assert_eq!(values_struct.data_type(), &DataType::Struct(fields));
1616
1617 let validity: Vec<bool> = {
1618 let null_buffer = values_struct.nulls().expect("should_have_nulls");
1619
1620 null_buffer.iter().collect()
1621 };
1622 assert_eq!(validity, &[true, true, true, false, true]);
1623 let values_number = values_struct.column(0).as_primitive::<Int32Type>();
1624 assert_eq!(values_number.values(), &[4, 4, 10, 6, 8]);
1625 let values_string = values_struct.column(1).as_string::<i32>();
1626 let values_string: Vec<_> = values_string.into_iter().collect();
1627 assert_eq!(
1628 &values_string,
1629 &[Some("d"), Some("d"), Some("z"), Some("world"), Some("x"),]
1630 );
1631 }
1632
1633 #[test]
1634 fn test_struct_empty() {
1635 let fields = Fields::from(vec![
1636 Field::new("number_col", DataType::Int32, false),
1637 Field::new("string_col", DataType::Utf8, false),
1638 ]);
1639 let a = {
1640 let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
1641 let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);
1642
1643 StructArray::try_new(
1644 fields.clone(),
1645 vec![Arc::new(number_col), Arc::new(string_col)],
1646 None,
1647 )
1648 .unwrap()
1649 };
1650 let v = interleave(&[&a], &[]).unwrap();
1651 assert!(v.is_empty());
1652 assert_eq!(v.data_type(), &DataType::Struct(fields));
1653 }
1654
1655 #[test]
1656 fn interleave_sparse_nulls() {
1657 let values = StringArray::from_iter_values((0..100).map(|x| x.to_string()));
1658 let keys = Int32Array::from_iter_values(0..10);
1659 let dict_a = DictionaryArray::new(keys, Arc::new(values));
1660 let values = StringArray::new_null(0);
1661 let keys = Int32Array::new_null(10);
1662 let dict_b = DictionaryArray::new(keys, Arc::new(values));
1663
1664 let indices = &[(0, 0), (0, 1), (0, 2), (1, 0)];
1665 let array = interleave(&[&dict_a, &dict_b], indices).unwrap();
1666
1667 let expected =
1668 DictionaryArray::<Int32Type>::from_iter(vec![Some("0"), Some("1"), Some("2"), None]);
1669 assert_eq!(array.as_ref(), &expected)
1670 }
1671
1672 #[test]
1673 fn test_interleave_views() {
1674 let values = StringArray::from_iter_values([
1675 "hello",
1676 "world_long_string_not_inlined",
1677 "foo",
1678 "bar",
1679 "baz",
1680 ]);
1681 let view_a = StringViewArray::from(&values);
1682
1683 let values = StringArray::from_iter_values([
1684 "test",
1685 "data",
1686 "more_long_string_not_inlined",
1687 "views",
1688 "here",
1689 ]);
1690 let view_b = StringViewArray::from(&values);
1691
1692 let indices = &[
1693 (0, 2), (1, 0), (0, 4), (1, 3), (0, 1), ];
1699
1700 let values = interleave(&[&view_a, &view_b], indices).unwrap();
1702 let result = values.as_string_view();
1703 assert_eq!(result.data_buffers().len(), 1);
1704
1705 let fallback = interleave_fallback(&[&view_a, &view_b], indices).unwrap();
1706 let fallback_result = fallback.as_string_view();
1707 assert_eq!(fallback_result.data_buffers().len(), 2);
1709
1710 let collected: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();
1712
1713 let fallback_collected: Vec<_> = fallback_result
1714 .iter()
1715 .map(|x| x.map(|s| s.to_string()))
1716 .collect();
1717
1718 assert_eq!(&collected, &fallback_collected);
1719
1720 assert_eq!(
1721 &collected,
1722 &[
1723 Some("foo".to_string()),
1724 Some("test".to_string()),
1725 Some("baz".to_string()),
1726 Some("views".to_string()),
1727 Some("world_long_string_not_inlined".to_string()),
1728 ]
1729 );
1730 }
1731
1732 #[test]
1733 fn test_interleave_views_with_nulls() {
1734 let values = StringArray::from_iter([
1735 Some("hello"),
1736 None,
1737 Some("foo_long_string_not_inlined"),
1738 Some("bar"),
1739 None,
1740 ]);
1741 let view_a = StringViewArray::from(&values);
1742
1743 let values = StringArray::from_iter([
1744 Some("test"),
1745 Some("data_long_string_not_inlined"),
1746 None,
1747 None,
1748 Some("here"),
1749 ]);
1750 let view_b = StringViewArray::from(&values);
1751
1752 let indices = &[
1753 (0, 1), (1, 2), (0, 2), (1, 3), (0, 4), ];
1759
1760 let values = interleave(&[&view_a, &view_b], indices).unwrap();
1762 let result = values.as_string_view();
1763 assert_eq!(result.data_buffers().len(), 1);
1764
1765 let fallback = interleave_fallback(&[&view_a, &view_b], indices).unwrap();
1766 let fallback_result = fallback.as_string_view();
1767
1768 let collected: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();
1770
1771 let fallback_collected: Vec<_> = fallback_result
1772 .iter()
1773 .map(|x| x.map(|s| s.to_string()))
1774 .collect();
1775
1776 assert_eq!(&collected, &fallback_collected);
1777
1778 assert_eq!(
1779 &collected,
1780 &[
1781 None,
1782 None,
1783 Some("foo_long_string_not_inlined".to_string()),
1784 None,
1785 None,
1786 ]
1787 );
1788 }
1789
1790 #[test]
1791 fn test_interleave_views_shared_buffers() {
1792 let long = |i: usize| format!("long_string_not_inlined_{i}");
1793 let base = StringViewArray::from_iter_values((0..10).map(long));
1794 assert_eq!(base.data_buffers().len(), 1);
1795
1796 let a = base.slice(0, 5);
1798 let b = base.slice(5, 5);
1799 let c = StringViewArray::try_new(base.views().clone(), base.data_buffers().to_vec(), None)
1801 .unwrap();
1802
1803 let indices = &[(0, 1), (1, 2), (2, 3), (0, 4), (1, 0), (2, 9)];
1804 let values = interleave(&[&a, &b, &c], indices).unwrap();
1805 let result = values.as_string_view();
1806 assert_eq!(result.data_buffers().len(), 1);
1807
1808 let expected: Vec<_> = [1, 7, 3, 4, 5, 9].into_iter().map(long).collect();
1809 let actual: Vec<_> = result.iter().map(|x| x.unwrap().to_string()).collect();
1810 assert_eq!(actual, expected);
1811
1812 let slices: Vec<_> = (0..40).map(|i| base.slice(i % 10, 1)).collect();
1814 let arrays: Vec<&dyn Array> = slices.iter().map(|x| x as &dyn Array).collect();
1815 let indices: Vec<_> = (0..40).rev().map(|i| (i, 0)).collect();
1816 let values = interleave(&arrays, &indices).unwrap();
1817 let result = values.as_string_view();
1818 assert_eq!(result.data_buffers().len(), 1);
1819
1820 let expected: Vec<_> = (0..40).rev().map(|i| long(i % 10)).collect();
1821 let actual: Vec<_> = result.iter().map(|x| x.unwrap().to_string()).collect();
1822 assert_eq!(actual, expected);
1823 }
1824
1825 #[test]
1826 fn test_interleave_views_multiple_buffers() {
1827 let str1 = "very_long_string_from_first_buffer".as_bytes();
1828 let str2 = "very_long_string_from_second_buffer".as_bytes();
1829 let buffer1 = str1.to_vec().into();
1830 let buffer2 = str2.to_vec().into();
1831
1832 let view1 = ByteView::new(str1.len() as u32, &str1[..4])
1833 .with_buffer_index(0)
1834 .with_offset(0)
1835 .as_u128();
1836 let view2 = ByteView::new(str2.len() as u32, &str2[..4])
1837 .with_buffer_index(1)
1838 .with_offset(0)
1839 .as_u128();
1840 let view_a =
1841 StringViewArray::try_new(vec![view1, view2].into(), vec![buffer1, buffer2], None)
1842 .unwrap();
1843
1844 let str3 = "another_very_long_string_buffer_three".as_bytes();
1845 let str4 = "different_long_string_in_buffer_four".as_bytes();
1846 let buffer3 = str3.to_vec().into();
1847 let buffer4 = str4.to_vec().into();
1848
1849 let view3 = ByteView::new(str3.len() as u32, &str3[..4])
1850 .with_buffer_index(0)
1851 .with_offset(0)
1852 .as_u128();
1853 let view4 = ByteView::new(str4.len() as u32, &str4[..4])
1854 .with_buffer_index(1)
1855 .with_offset(0)
1856 .as_u128();
1857 let view_b =
1858 StringViewArray::try_new(vec![view3, view4].into(), vec![buffer3, buffer4], None)
1859 .unwrap();
1860
1861 let indices = &[
1862 (0, 0), (1, 0), (0, 1), (1, 1), (0, 0), (1, 1), ];
1869
1870 let values = interleave(&[&view_a, &view_b], indices).unwrap();
1872 let result = values.as_string_view();
1873
1874 assert_eq!(
1875 result.data_buffers().len(),
1876 4,
1877 "Expected four buffers (two from each input array)"
1878 );
1879
1880 let result_strings: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();
1881 assert_eq!(
1882 result_strings,
1883 vec![
1884 Some("very_long_string_from_first_buffer".to_string()),
1885 Some("another_very_long_string_buffer_three".to_string()),
1886 Some("very_long_string_from_second_buffer".to_string()),
1887 Some("different_long_string_in_buffer_four".to_string()),
1888 Some("very_long_string_from_first_buffer".to_string()),
1889 Some("different_long_string_in_buffer_four".to_string()),
1890 ]
1891 );
1892
1893 let views = result.views();
1894 let buffer_indices: Vec<_> = views
1895 .iter()
1896 .map(|raw_view| ByteView::from(*raw_view).buffer_index)
1897 .collect();
1898
1899 assert_eq!(
1900 buffer_indices,
1901 vec![
1902 0, 1, 2, 3, 0, 3, ]
1909 );
1910 }
1911
1912 #[test]
1913 fn test_interleave_run_end_encoded_primitive() {
1914 let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
1915 builder.extend([1, 1, 2, 2, 2, 3].into_iter().map(Some));
1916 let a = builder.finish();
1917
1918 let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
1919 builder.extend([4, 5, 5, 6, 6, 6].into_iter().map(Some));
1920 let b = builder.finish();
1921
1922 let indices = &[(0, 1), (1, 0), (0, 4), (1, 2), (0, 5)];
1923 let result = interleave(&[&a, &b], indices).unwrap();
1924
1925 assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
1927
1928 let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();
1930
1931 let expected = vec![1, 4, 2, 5, 3];
1933 let mut actual = Vec::new();
1934 for i in 0..result_run_array.len() {
1935 let physical_idx = result_run_array.get_physical_index(i);
1936 let value = result_run_array
1937 .values()
1938 .as_primitive::<Int32Type>()
1939 .value(physical_idx);
1940 actual.push(value);
1941 }
1942 assert_eq!(actual, expected);
1943 }
1944
1945 #[test]
1946 fn test_interleave_run_end_encoded_sliced() {
1947 let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
1948 builder.extend([1, 1, 2, 2, 2, 3].into_iter().map(Some));
1949 let a = builder.finish();
1950 let a = a.slice(2, 3); let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
1953 builder.extend([4, 5, 5, 6, 6, 6].into_iter().map(Some));
1954 let b = builder.finish();
1955 let b = b.slice(1, 3); let indices = &[(0, 1), (1, 0), (0, 2), (1, 1), (1, 2)];
1958 let result = interleave(&[&a, &b], indices).unwrap();
1959
1960 let result = result.as_run::<Int32Type>();
1961 let result = result.downcast::<Int32Array>().unwrap();
1962
1963 let expected = vec![2, 5, 2, 5, 6];
1964 let actual = result.into_iter().flatten().collect::<Vec<_>>();
1965 assert_eq!(actual, expected);
1966 }
1967
1968 #[test]
1969 fn test_interleave_run_end_encoded_string() {
1970 let a: Int32RunArray = vec!["hello", "hello", "world", "world", "foo"]
1971 .into_iter()
1972 .collect();
1973 let b: Int32RunArray = vec!["bar", "baz", "baz", "qux"].into_iter().collect();
1974
1975 let indices = &[(0, 0), (1, 1), (0, 3), (1, 3), (0, 4)];
1976 let result = interleave(&[&a, &b], indices).unwrap();
1977
1978 assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
1980
1981 let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();
1983
1984 let expected = vec!["hello", "baz", "world", "qux", "foo"];
1986 let mut actual = Vec::new();
1987 for i in 0..result_run_array.len() {
1988 let physical_idx = result_run_array.get_physical_index(i);
1989 let value = result_run_array
1990 .values()
1991 .as_string::<i32>()
1992 .value(physical_idx);
1993 actual.push(value);
1994 }
1995 assert_eq!(actual, expected);
1996 }
1997
1998 #[test]
1999 fn test_interleave_run_end_encoded_with_nulls() {
2000 let a: Int32RunArray = vec![Some("a"), Some("a"), None, None, Some("b")]
2001 .into_iter()
2002 .collect();
2003 let b: Int32RunArray = vec![None, Some("c"), Some("c"), Some("d")]
2004 .into_iter()
2005 .collect();
2006
2007 let indices = &[(0, 1), (1, 0), (0, 2), (1, 3), (0, 4)];
2008 let result = interleave(&[&a, &b], indices).unwrap();
2009
2010 assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
2012
2013 let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();
2015
2016 let expected = vec![Some("a"), None, None, Some("d"), Some("b")];
2018 let mut actual = Vec::new();
2019 for i in 0..result_run_array.len() {
2020 let physical_idx = result_run_array.get_physical_index(i);
2021 if result_run_array.values().is_null(physical_idx) {
2022 actual.push(None);
2023 } else {
2024 let value = result_run_array
2025 .values()
2026 .as_string::<i32>()
2027 .value(physical_idx);
2028 actual.push(Some(value));
2029 }
2030 }
2031 assert_eq!(actual, expected);
2032 }
2033
2034 #[test]
2035 fn test_interleave_run_end_encoded_different_run_types() {
2036 let mut builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2037 builder.extend([1, 1, 2, 3, 3].into_iter().map(Some));
2038 let a = builder.finish();
2039
2040 let mut builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2041 builder.extend([4, 5, 5, 6].into_iter().map(Some));
2042 let b = builder.finish();
2043
2044 let indices = &[(0, 0), (1, 1), (0, 3), (1, 3)];
2045 let result = interleave(&[&a, &b], indices).unwrap();
2046
2047 assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
2049
2050 let result_run_array: &RunArray<Int16Type> = result.as_any().downcast_ref().unwrap();
2052
2053 let expected = vec![1, 5, 3, 6];
2055 let mut actual = Vec::new();
2056 for i in 0..result_run_array.len() {
2057 let physical_idx = result_run_array.get_physical_index(i);
2058 let value = result_run_array
2059 .values()
2060 .as_primitive::<Int32Type>()
2061 .value(physical_idx);
2062 actual.push(value);
2063 }
2064 assert_eq!(actual, expected);
2065 }
2066
2067 #[test]
2068 fn test_interleave_run_end_encoded_mixed_run_lengths() {
2069 let mut builder = PrimitiveRunBuilder::<Int64Type, Int32Type>::new();
2070 builder.extend([1, 2, 2, 2, 2, 3, 3, 4].into_iter().map(Some));
2071 let a = builder.finish();
2072
2073 let mut builder = PrimitiveRunBuilder::<Int64Type, Int32Type>::new();
2074 builder.extend([5, 5, 5, 6, 7, 7, 8, 8].into_iter().map(Some));
2075 let b = builder.finish();
2076
2077 let indices = &[
2078 (0, 0), (1, 2), (0, 3), (1, 3), (0, 6), (1, 6), (0, 7), (1, 4), ];
2087 let result = interleave(&[&a, &b], indices).unwrap();
2088
2089 assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
2091
2092 let result_run_array: &RunArray<Int64Type> = result.as_any().downcast_ref().unwrap();
2094
2095 let expected = vec![1, 5, 2, 6, 3, 8, 4, 7];
2097 let mut actual = Vec::new();
2098 for i in 0..result_run_array.len() {
2099 let physical_idx = result_run_array.get_physical_index(i);
2100 let value = result_run_array
2101 .values()
2102 .as_primitive::<Int32Type>()
2103 .value(physical_idx);
2104 actual.push(value);
2105 }
2106 assert_eq!(actual, expected);
2107 }
2108
2109 #[test]
2110 fn test_interleave_run_end_encoded_empty_runs() {
2111 let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
2112 builder.extend(std::iter::once(Some(1)));
2113 let a = builder.finish();
2114
2115 let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
2116 builder.extend([2, 2, 2].into_iter().map(Some));
2117 let b = builder.finish();
2118
2119 let indices = &[(0, 0), (1, 1), (1, 2)];
2120 let result = interleave(&[&a, &b], indices).unwrap();
2121
2122 assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
2124
2125 let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();
2127
2128 let expected = vec![1, 2, 2];
2130 let mut actual = Vec::new();
2131 for i in 0..result_run_array.len() {
2132 let physical_idx = result_run_array.get_physical_index(i);
2133 let value = result_run_array
2134 .values()
2135 .as_primitive::<Int32Type>()
2136 .value(physical_idx);
2137 actual.push(value);
2138 }
2139 assert_eq!(actual, expected);
2140 }
2141
2142 #[test]
2143 fn test_struct_no_fields() {
2144 let fields = Fields::empty();
2145 let a = StructArray::try_new_with_length(fields.clone(), vec![], None, 10).unwrap();
2146 let v = interleave(&[&a], &[(0, 0)]).unwrap();
2147 assert_eq!(v.len(), 1);
2148 assert_eq!(v.data_type(), &DataType::Struct(fields));
2149 }
2150
2151 #[test]
2152 fn test_interleave_fallback_dictionary_with_nulls() {
2153 let input_1_keys = Int32Array::from_iter([Some(0), None, Some(1)]);
2154 let input_1_values = StringArray::from_iter_values(["foo", "bar"]);
2155 let dict_a = DictionaryArray::new(input_1_keys, Arc::new(input_1_values));
2156
2157 let input_2_keys = Int32Array::from_iter([Some(0), Some(1), None]);
2158 let input_2_values = StringArray::from_iter_values(["baz", "qux"]);
2159 let dict_b = DictionaryArray::new(input_2_keys, Arc::new(input_2_values));
2160
2161 let indices = vec![
2162 (0, 0), (0, 1), (1, 0), (1, 2), (0, 2), (1, 1), ];
2169
2170 let result =
2171 interleave_fallback_dictionary::<Int32Type>(&[&dict_a, &dict_b], &indices).unwrap();
2172 let dict_result = result.as_dictionary::<Int32Type>();
2173
2174 let string_result = dict_result.downcast_dict::<StringArray>().unwrap();
2175 let collected: Vec<_> = string_result.into_iter().collect();
2176 assert_eq!(
2177 collected,
2178 vec![
2179 Some("foo"),
2180 None,
2181 Some("baz"),
2182 None,
2183 Some("bar"),
2184 Some("qux")
2185 ]
2186 );
2187 }
2188
2189 #[test]
2190 fn test_interleave_string_view_dictionary_overflow_returns_err() {
2191 let values_a: StringViewArray = (0..200).map(|i| Some(format!("a{i}"))).collect();
2194 let keys_a = UInt8Array::from_iter_values(0..200);
2195 let dict_a = DictionaryArray::<UInt8Type>::new(keys_a, Arc::new(values_a));
2196
2197 let values_b: StringViewArray = (0..200).map(|i| Some(format!("b{i}"))).collect();
2198 let keys_b = UInt8Array::from_iter_values(0..200);
2199 let dict_b = DictionaryArray::<UInt8Type>::new(keys_b, Arc::new(values_b));
2200
2201 let indices: Vec<_> = (0..200).flat_map(|i| [(0, i), (1, i)]).collect();
2202
2203 let err = interleave(&[&dict_a, &dict_b], &indices).unwrap_err();
2204 assert!(matches!(err, ArrowError::DictionaryKeyOverflowError));
2205 }
2206
2207 #[test]
2208 fn test_interleave_nested_dictionary_overflow_returns_err() {
2209 let field = Arc::new(arrow_schema::Field::new(
2211 "item",
2212 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8View)),
2213 false,
2214 ));
2215
2216 let values_a: StringViewArray = (0..200).map(|i| Some(format!("a{i}"))).collect();
2217 let keys_a = UInt8Array::from_iter_values(0..200);
2218 let dict_a = DictionaryArray::<UInt8Type>::new(keys_a, Arc::new(values_a));
2219 let list_a = FixedSizeListArray::new(field.clone(), 1, Arc::new(dict_a), None);
2220
2221 let values_b: StringViewArray = (0..200).map(|i| Some(format!("b{i}"))).collect();
2222 let keys_b = UInt8Array::from_iter_values(0..200);
2223 let dict_b = DictionaryArray::<UInt8Type>::new(keys_b, Arc::new(values_b));
2224 let list_b = FixedSizeListArray::new(field, 1, Arc::new(dict_b), None);
2225
2226 let indices: Vec<_> = (0..200).flat_map(|i| [(0, i), (1, i)]).collect();
2227
2228 let err = interleave(&[&list_a, &list_b], &indices).unwrap_err();
2229 assert!(matches!(err, ArrowError::DictionaryKeyOverflowError));
2230 }
2231
2232 #[test]
2233 #[cfg_attr(miri, ignore)] fn test_interleave_bytes_offset_overflow() {
2235 let indices: Vec<(usize, usize)> = vec![(0, 0); (i32::MAX >> 4) as usize];
2236 let text = ('a'..='z').collect::<String>();
2237 let values = StringArray::from(vec![Some(text)]);
2238 assert!(matches!(
2239 interleave(&[&values], &indices),
2240 Err(ArrowError::OffsetOverflowError(_))
2241 ));
2242 }
2243
2244 #[test]
2245 #[cfg_attr(miri, ignore)] fn test_interleave_list_offset_overflow() {
2247 let mut builder = GenericListBuilder::<i32, _>::new(Int32Builder::new());
2249 for i in 0..32 {
2250 builder.values().append_value(i);
2251 }
2252 builder.append(true);
2253 let list = builder.finish();
2254
2255 let indices: Vec<(usize, usize)> = vec![(0, 0); (i32::MAX as usize / 32) + 1];
2257 assert!(matches!(
2258 interleave(&[&list], &indices),
2259 Err(ArrowError::OffsetOverflowError(_))
2260 ));
2261 }
2262
2263 #[test]
2264 fn test_interleave_list_view() {
2265 let field = Arc::new(Field::new_list_field(DataType::Int64, false));
2274
2275 let lv_a = ListViewArray::new(
2276 Arc::clone(&field),
2277 ScalarBuffer::from(vec![0i32, 2]),
2278 ScalarBuffer::from(vec![2i32, 1]),
2279 Arc::new(Int64Array::from(vec![1_i64, 2, 3])),
2280 None,
2281 );
2282 let lv_b = ListViewArray::new(
2283 field,
2284 ScalarBuffer::from(vec![0i32]),
2285 ScalarBuffer::from(vec![3i32]),
2286 Arc::new(Int64Array::from(vec![4_i64, 5, 6])),
2287 None,
2288 );
2289
2290 let result = interleave(
2291 &[&lv_a as &dyn Array, &lv_b as &dyn Array],
2292 &[(0, 0), (1, 0), (0, 1)],
2293 )
2294 .unwrap();
2295
2296 result
2297 .to_data()
2298 .validate_full()
2299 .expect("interleaved ListViewArray must be internally consistent");
2300
2301 let result_lv = result.as_list_view::<i32>();
2302 assert_eq!(result_lv.len(), 3);
2303 assert_eq!(
2304 result_lv.value(0).as_primitive::<Int64Type>().values(),
2305 &[1, 2]
2306 );
2307 assert_eq!(
2308 result_lv.value(1).as_primitive::<Int64Type>().values(),
2309 &[4, 5, 6]
2310 );
2311 assert_eq!(
2312 result_lv.value(2).as_primitive::<Int64Type>().values(),
2313 &[3]
2314 );
2315 }
2316
2317 #[test]
2318 fn test_interleave_fixed_size_list() {
2319 let field = Arc::new(Field::new("item", DataType::Int32, false));
2321 let a = FixedSizeListArray::new(
2322 field.clone(),
2323 2,
2324 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])),
2325 None,
2326 );
2327 let b = FixedSizeListArray::new(
2329 field.clone(),
2330 2,
2331 Arc::new(Int32Array::from(vec![7, 8, 9, 10])),
2332 None,
2333 );
2334
2335 let result = interleave(&[&a, &b], &[(0, 2), (1, 0), (0, 0), (1, 1), (0, 1)]).unwrap();
2336 let result = result.as_fixed_size_list();
2337 assert_eq!(result.len(), 5);
2338 assert_eq!(result.value_length(), 2);
2339
2340 let values = result.values().as_primitive::<Int32Type>();
2341 assert_eq!(values.values(), &[5, 6, 7, 8, 1, 2, 9, 10, 3, 4]);
2343 }
2344
2345 #[test]
2346 fn test_interleave_zero_sized_fixed_size_list() {
2347 let input = FixedSizeListArray::try_new_with_length(
2348 Field::new_list_field(DataType::Int32, true).into(),
2349 0,
2350 Arc::new(Int32Array::new_null(0)),
2351 None,
2352 3,
2353 )
2354 .unwrap();
2355
2356 let indices = [(0, 2), (0, 0)];
2357 let result = interleave(&[&input], &indices).unwrap();
2358
2359 assert_eq!(result.len(), 2);
2360 }
2361
2362 #[test]
2363 fn test_interleave_fixed_size_list_with_nulls() {
2364 let field = Arc::new(Field::new("item", DataType::Int32, true));
2365 let a = FixedSizeListArray::new(
2367 field.clone(),
2368 2,
2369 Arc::new(Int32Array::from(vec![1, 2, 0, 0, 5, 6])),
2370 Some(NullBuffer::from(&[true, false, true])),
2371 );
2372 let b = FixedSizeListArray::new(
2374 field.clone(),
2375 2,
2376 Arc::new(Int32Array::from(vec![0, 0, 9, 10])),
2377 Some(NullBuffer::from(&[false, true])),
2378 );
2379
2380 let result = interleave(&[&a, &b], &[(0, 0), (0, 1), (1, 0), (1, 1), (0, 2)]).unwrap();
2381 let result = result.as_fixed_size_list();
2382 assert_eq!(result.len(), 5);
2383
2384 let validity: Vec<bool> = result.nulls().unwrap().iter().collect();
2385 assert_eq!(validity, &[true, false, false, true, true]);
2386 }
2387
2388 fn run_fsl_string_child_test(
2389 child_data_type: DataType,
2390 create_child: impl Fn(Vec<&str>) -> ArrayRef,
2391 extract_values: impl Fn(&ArrayRef) -> Vec<String>,
2392 ) {
2393 let field = Arc::new(Field::new("item", child_data_type, false));
2394
2395 let a = FixedSizeListArray::new(
2397 field.clone(),
2398 2,
2399 create_child(vec!["a", "b", "c", "d"]),
2400 None,
2401 );
2402 let b = FixedSizeListArray::new(
2404 field.clone(),
2405 2,
2406 create_child(vec!["x", "y", "z", "w"]),
2407 None,
2408 );
2409
2410 let result = interleave(&[&a, &b], &[(0, 1), (1, 0), (0, 0)]).unwrap();
2411 let result = result.as_fixed_size_list();
2412 assert_eq!(result.len(), 3);
2413 assert_eq!(result.value_length(), 2);
2414
2415 let values = extract_values(result.values());
2417 assert_eq!(values, vec!["c", "d", "x", "y", "a", "b"]);
2418 }
2419
2420 #[test]
2421 fn test_interleave_fixed_size_list_string_child() {
2422 run_fsl_string_child_test(
2424 DataType::Utf8,
2425 |v| Arc::new(StringArray::from(v)),
2426 |a| {
2427 a.as_string::<i32>()
2428 .iter()
2429 .map(|s| s.unwrap().to_owned())
2430 .collect()
2431 },
2432 );
2433 }
2434
2435 #[test]
2436 fn test_interleave_fixed_size_list_string_view_child() {
2437 run_fsl_string_child_test(
2439 DataType::Utf8View,
2440 |v| Arc::new(StringViewArray::from(v)),
2441 |a| {
2442 a.as_string_view()
2443 .iter()
2444 .map(|s| s.unwrap().to_owned())
2445 .collect()
2446 },
2447 );
2448 }
2449
2450 #[test]
2451 fn test_interleave_map() {
2452 use arrow_array::builder::MapBuilder;
2453 use arrow_array::builder::StringBuilder;
2454
2455 let mut a_builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
2457 a_builder.keys().append_value("k1");
2458 a_builder.values().append_value(1);
2459 a_builder.keys().append_value("k2");
2460 a_builder.values().append_value(2);
2461 a_builder.append(true).unwrap();
2462 a_builder.keys().append_value("k3");
2463 a_builder.values().append_value(3);
2464 a_builder.append(true).unwrap();
2465 let a = a_builder.finish();
2466
2467 let mut b_builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
2469 b_builder.keys().append_value("k4");
2470 b_builder.values().append_value(4);
2471 b_builder.append(true).unwrap();
2472 b_builder.keys().append_value("k5");
2473 b_builder.values().append_value(5);
2474 b_builder.keys().append_value("k6");
2475 b_builder.values().append_value(6);
2476 b_builder.keys().append_value("k7");
2477 b_builder.values().append_value(7);
2478 b_builder.append(true).unwrap();
2479 let b = b_builder.finish();
2480
2481 let result = interleave(&[&a, &b], &[(1, 0), (0, 0), (0, 1), (1, 1)]).unwrap();
2482 let result = result.as_map();
2483 assert_eq!(result.len(), 4);
2484
2485 let row0 = result.value(0);
2487 assert_eq!(row0.len(), 1);
2488 assert_eq!(row0.column(0).as_string::<i32>().value(0), "k4");
2489 assert_eq!(row0.column(1).as_primitive::<Int32Type>().value(0), 4);
2490
2491 let row1 = result.value(1);
2493 assert_eq!(row1.len(), 2);
2494
2495 let row2 = result.value(2);
2497 assert_eq!(row2.len(), 1);
2498 assert_eq!(row2.column(0).as_string::<i32>().value(0), "k3");
2499
2500 let row3 = result.value(3);
2502 assert_eq!(row3.len(), 3);
2503 }
2504
2505 #[test]
2506 fn test_interleave_map_with_nulls() {
2507 use arrow_array::builder::MapBuilder;
2508 use arrow_array::builder::StringBuilder;
2509
2510 let mut a_builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
2512 a_builder.keys().append_value("k1");
2513 a_builder.values().append_value(1);
2514 a_builder.append(true).unwrap();
2515 a_builder.append(false).unwrap();
2516 let a = a_builder.finish();
2517
2518 let mut b_builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
2520 b_builder.append(false).unwrap();
2521 b_builder.keys().append_value("k2");
2522 b_builder.values().append_value(2);
2523 b_builder.append(true).unwrap();
2524 let b = b_builder.finish();
2525
2526 let result = interleave(&[&a, &b], &[(0, 0), (1, 0), (0, 1), (1, 1)]).unwrap();
2527 let result = result.as_map();
2528 assert_eq!(result.len(), 4);
2529
2530 let validity: Vec<bool> = result.nulls().unwrap().iter().collect();
2531 assert_eq!(validity, &[true, false, false, true]);
2532
2533 assert_eq!(result.value(0).len(), 1);
2534 assert_eq!(result.value(3).len(), 1);
2535 }
2536}