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::{ArrowNativeType, BooleanBuffer, MutableBuffer, NullBuffer, OffsetBuffer};
30use arrow_data::ByteView;
31use arrow_data::transform::MutableArrayData;
32use arrow_schema::{ArrowError, DataType, FieldRef, Fields};
33use std::sync::Arc;
34
35macro_rules! primitive_helper {
36 ($t:ty, $values:ident, $indices:ident, $data_type:ident) => {
37 interleave_primitive::<$t>($values, $indices, $data_type)
38 };
39}
40
41macro_rules! dict_helper {
42 ($t:ty, $values:expr, $indices:expr) => {
43 interleave_dictionaries::<$t>($values, $indices)
44 };
45}
46
47pub fn interleave(
75 values: &[&dyn Array],
76 indices: &[(usize, usize)],
77) -> Result<ArrayRef, ArrowError> {
78 if values.is_empty() {
79 return Err(ArrowError::InvalidArgumentError(
80 "interleave requires input of at least one array".to_string(),
81 ));
82 }
83 let data_type = values[0].data_type();
84
85 for array in values.iter().skip(1) {
86 if array.data_type() != data_type {
87 return Err(ArrowError::InvalidArgumentError(format!(
88 "It is not possible to interleave arrays of different data types ({} and {})",
89 data_type,
90 array.data_type()
91 )));
92 }
93 }
94
95 if indices.is_empty() {
96 return Ok(new_empty_array(data_type));
97 }
98
99 downcast_primitive! {
100 data_type => (primitive_helper, values, indices, data_type),
101 DataType::Utf8 => interleave_bytes::<Utf8Type>(values, indices),
102 DataType::LargeUtf8 => interleave_bytes::<LargeUtf8Type>(values, indices),
103 DataType::Binary => interleave_bytes::<BinaryType>(values, indices),
104 DataType::LargeBinary => interleave_bytes::<LargeBinaryType>(values, indices),
105 DataType::BinaryView => interleave_views::<BinaryViewType>(values, indices),
106 DataType::Utf8View => interleave_views::<StringViewType>(values, indices),
107 DataType::Dictionary(k, _) => downcast_integer! {
108 k.as_ref() => (dict_helper, values, indices),
109 _ => unreachable!("illegal dictionary key type {k}")
110 },
111 DataType::Struct(fields) => interleave_struct(fields, values, indices),
112 DataType::List(field) => interleave_list::<i32>(values, indices, field),
113 DataType::LargeList(field) => interleave_list::<i64>(values, indices, field),
114 DataType::FixedSizeList(field, size) => interleave_fixed_size_list(values, indices, field, *size),
115 DataType::Map(field, ordered) => interleave_map(values, indices, field, *ordered),
116 DataType::RunEndEncoded(r, _) => match r.data_type() {
117 DataType::Int16 => interleave_run_end::<Int16Type>(values, indices),
118 DataType::Int32 => interleave_run_end::<Int32Type>(values, indices),
119 DataType::Int64 => interleave_run_end::<Int64Type>(values, indices),
120 t => unreachable!("illegal run-end type {t}"),
121 },
122 DataType::ListView(field) => interleave_list_view::<i32>(values, indices, field),
123 DataType::LargeListView(field) => interleave_list_view::<i64>(values, indices, field),
124 _ => interleave_fallback(values, indices)
125 }
126}
127
128struct Interleave<'a, T> {
132 arrays: Vec<&'a T>,
134 nulls: Option<NullBuffer>,
136}
137
138impl<'a, T: Array + 'static> Interleave<'a, T> {
139 fn new(values: &[&'a dyn Array], indices: &'a [(usize, usize)]) -> Self {
140 let mut has_nulls = false;
141 let arrays: Vec<&T> = values
142 .iter()
143 .map(|x| {
144 has_nulls = has_nulls || x.null_count() != 0;
145 x.as_any().downcast_ref().unwrap()
146 })
147 .collect();
148
149 let nulls = match has_nulls {
150 true => {
151 let nulls = BooleanBuffer::collect_bool(indices.len(), |i| {
152 let (a, b) = indices[i];
153 arrays[a].is_valid(b)
154 });
155 Some(nulls.into())
156 }
157 false => None,
158 };
159
160 Self { arrays, nulls }
161 }
162}
163
164fn interleave_primitive<T: ArrowPrimitiveType>(
165 values: &[&dyn Array],
166 indices: &[(usize, usize)],
167 data_type: &DataType,
168) -> Result<ArrayRef, ArrowError> {
169 let interleaved = Interleave::<'_, PrimitiveArray<T>>::new(values, indices);
170 let arrays = &interleaved.arrays;
171 let len = indices.len();
172
173 let mut output = Vec::with_capacity(len);
174 let dst: *mut T::Native = output.as_mut_ptr();
175 let mut base = 0;
176
177 let chunks = indices.chunks_exact(8);
180 let remainder = chunks.remainder();
181 for chunk in chunks {
182 let v0 = arrays[chunk[0].0].value(chunk[0].1);
183 let v1 = arrays[chunk[1].0].value(chunk[1].1);
184 let v2 = arrays[chunk[2].0].value(chunk[2].1);
185 let v3 = arrays[chunk[3].0].value(chunk[3].1);
186 let v4 = arrays[chunk[4].0].value(chunk[4].1);
187 let v5 = arrays[chunk[5].0].value(chunk[5].1);
188 let v6 = arrays[chunk[6].0].value(chunk[6].1);
189 let v7 = arrays[chunk[7].0].value(chunk[7].1);
190
191 debug_assert!(base + 7 < len);
193 unsafe {
194 dst.add(base).write(v0);
195 dst.add(base + 1).write(v1);
196 dst.add(base + 2).write(v2);
197 dst.add(base + 3).write(v3);
198 dst.add(base + 4).write(v4);
199 dst.add(base + 5).write(v5);
200 dst.add(base + 6).write(v6);
201 dst.add(base + 7).write(v7);
202 }
203 base += 8;
204 }
205
206 for idx in remainder {
207 debug_assert!(base < len);
209 unsafe { dst.add(base).write(arrays[idx.0].value(idx.1)) };
210 base += 1;
211 }
212
213 debug_assert!(base == len);
215 unsafe { output.set_len(len) };
216
217 let array = PrimitiveArray::<T>::try_new(output.into(), interleaved.nulls)?;
218 Ok(Arc::new(array.with_data_type(data_type.clone())))
219}
220
221fn interleave_bytes<T: ByteArrayType>(
222 values: &[&dyn Array],
223 indices: &[(usize, usize)],
224) -> Result<ArrayRef, ArrowError> {
225 let interleaved = Interleave::<'_, GenericByteArray<T>>::new(values, indices);
226
227 let mut capacity = 0;
228 let mut offsets = Vec::with_capacity(indices.len() + 1);
229 offsets.push(T::Offset::from_usize(0).unwrap());
230 for (a, b) in indices {
231 let o = interleaved.arrays[*a].value_offsets();
232 let element_len = o[*b + 1].as_usize() - o[*b].as_usize();
233 capacity += element_len;
234 offsets.push(
235 T::Offset::from_usize(capacity)
236 .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?,
237 );
238 }
239
240 let mut values = Vec::with_capacity(capacity);
241 for (a, b) in indices {
242 values.extend_from_slice(interleaved.arrays[*a].value(*b).as_ref());
243 }
244
245 let array = unsafe {
247 let offsets = OffsetBuffer::new_unchecked(offsets.into());
248 GenericByteArray::<T>::new_unchecked(offsets, values.into(), interleaved.nulls)
249 };
250 Ok(Arc::new(array))
251}
252
253fn interleave_dictionaries<K: ArrowDictionaryKeyType>(
254 arrays: &[&dyn Array],
255 indices: &[(usize, usize)],
256) -> Result<ArrayRef, ArrowError> {
257 let dictionaries: Vec<_> = arrays.iter().map(|x| x.as_dictionary::<K>()).collect();
258 let (should_merge, has_overflow) =
259 should_merge_dictionary_values::<K>(&dictionaries, indices.len());
260 if !should_merge {
261 return if has_overflow {
262 interleave_fallback(arrays, indices)
263 } else {
264 interleave_fallback_dictionary::<K>(&dictionaries, indices)
265 };
266 }
267
268 let masks: Vec<_> = dictionaries
269 .iter()
270 .enumerate()
271 .map(|(a_idx, dictionary)| {
272 let mut key_mask = BooleanBufferBuilder::new_from_buffer(
273 MutableBuffer::new_null(dictionary.len()),
274 dictionary.len(),
275 );
276
277 for (_, key_idx) in indices.iter().filter(|(a, _)| *a == a_idx) {
278 key_mask.set_bit(*key_idx, true);
279 }
280 key_mask.finish()
281 })
282 .collect();
283
284 let merged = merge_dictionary_values(&dictionaries, Some(&masks))?;
285
286 let mut keys = PrimitiveBuilder::<K>::with_capacity(indices.len());
288 for (a, b) in indices {
289 let old_keys: &PrimitiveArray<K> = dictionaries[*a].keys();
290 match old_keys.is_valid(*b) {
291 true => {
292 let old_key = old_keys.values()[*b];
293 keys.append_value(merged.key_mappings[*a][old_key.as_usize()])
294 }
295 false => keys.append_null(),
296 }
297 }
298 let array = unsafe { DictionaryArray::new_unchecked(keys.finish(), merged.values) };
299 Ok(Arc::new(array))
300}
301
302fn interleave_views<T: ByteViewType>(
303 values: &[&dyn Array],
304 indices: &[(usize, usize)],
305) -> Result<ArrayRef, ArrowError> {
306 let interleaved = Interleave::<'_, GenericByteViewArray<T>>::new(values, indices);
307 let mut buffers = Vec::new();
308
309 let mut offsets = Vec::with_capacity(interleaved.arrays.len() + 1);
311 offsets.push(0);
312 let mut total_buffers = 0;
313 for a in interleaved.arrays.iter() {
314 total_buffers += a.data_buffers().len();
315 offsets.push(total_buffers);
316 }
317
318 let mut buffer_to_new_index = vec![None; total_buffers];
320
321 let views: Vec<u128> = indices
322 .iter()
323 .map(|(array_idx, value_idx)| {
324 let array = interleaved.arrays[*array_idx];
325 let view = array.views().get(*value_idx).unwrap();
326 let view_len = *view as u32;
327 if view_len <= 12 {
328 return *view;
329 }
330 let view = ByteView::from(*view);
332 let buffer_to_new_idx = offsets[*array_idx] + view.buffer_index as usize;
333 let new_buffer_idx: u32 =
334 *buffer_to_new_index[buffer_to_new_idx].get_or_insert_with(|| {
335 buffers.push(array.data_buffers()[view.buffer_index as usize].clone());
336 (buffers.len() - 1) as u32
337 });
338 view.with_buffer_index(new_buffer_idx).as_u128()
339 })
340 .collect();
341
342 let array = unsafe {
343 GenericByteViewArray::<T>::new_unchecked(views.into(), buffers, interleaved.nulls)
344 };
345 Ok(Arc::new(array))
346}
347
348fn interleave_struct(
349 fields: &Fields,
350 values: &[&dyn Array],
351 indices: &[(usize, usize)],
352) -> Result<ArrayRef, ArrowError> {
353 let interleaved = Interleave::<'_, StructArray>::new(values, indices);
354
355 if fields.is_empty() {
356 let array = StructArray::try_new_with_length(
357 fields.clone(),
358 vec![],
359 interleaved.nulls,
360 indices.len(),
361 )?;
362 return Ok(Arc::new(array));
363 }
364
365 let struct_fields_array: Result<Vec<_>, _> = (0..fields.len())
366 .map(|i| {
367 let field_values: Vec<&dyn Array> = interleaved
368 .arrays
369 .iter()
370 .map(|x| x.column(i).as_ref())
371 .collect();
372 interleave(&field_values, indices)
373 })
374 .collect();
375
376 let struct_array =
377 StructArray::try_new(fields.clone(), struct_fields_array?, interleaved.nulls)?;
378 Ok(Arc::new(struct_array))
379}
380
381fn interleave_list_like_primitive_child<L: ListLikeArray, T: ArrowPrimitiveType>(
382 interleaved: &Interleave<'_, L>,
383 indices: &[(usize, usize)],
384 capacity: usize,
385 data_type: &DataType,
386) -> ArrayRef {
387 let child_arrays: Vec<&PrimitiveArray<T>> = interleaved
388 .arrays
389 .iter()
390 .map(|list| list.values().as_primitive::<T>())
391 .collect();
392
393 let has_child_nulls = child_arrays.iter().any(|a| a.null_count() > 0);
394
395 let mut values: Vec<T::Native> = Vec::with_capacity(capacity);
397 for &(array, row) in indices {
398 let range = interleaved.arrays[array].element_range(row);
399 if !range.is_empty() {
400 values.extend_from_slice(&child_arrays[array].values()[range]);
401 }
402 }
403
404 let nulls = if has_child_nulls {
408 let null_byte_len = bit_util::ceil(capacity, 8);
409 let mut output_null_buf = MutableBuffer::from_len_zeroed(null_byte_len);
410
411 let mut offset_write = 0;
412 let mut output_null_count = 0usize;
413 for &(array, row) in indices {
414 let range = interleaved.arrays[array].element_range(row);
415 let len = range.len();
416 if len > 0 {
417 match child_arrays[array].nulls() {
418 Some(null_buffer) => {
419 output_null_count += set_bits(
420 output_null_buf.as_slice_mut(),
421 null_buffer.validity(),
422 offset_write,
423 null_buffer.offset() + range.start,
424 len,
425 );
426 }
427 None => {
428 let buf = output_null_buf.as_slice_mut();
430 (offset_write..offset_write + len).for_each(|i| bit_util::set_bit(buf, i));
431 }
432 }
433 }
434 offset_write += len;
435 }
436
437 if output_null_count > 0 {
438 let bool_buf = BooleanBuffer::new(output_null_buf.into(), 0, capacity);
439 Some(unsafe { NullBuffer::new_unchecked(bool_buf, output_null_count) })
441 } else {
442 None
443 }
444 } else {
445 None
446 };
447
448 Arc::new(PrimitiveArray::<T>::new(values.into(), nulls).with_data_type(data_type.clone()))
449}
450
451fn interleave_list_like_child<L: ListLikeArray>(
453 interleaved: &Interleave<'_, L>,
454 indices: &[(usize, usize)],
455 capacity: usize,
456) -> Result<ArrayRef, ArrowError> {
457 let mut child_indices = Vec::with_capacity(capacity);
458 for &(array, row) in indices {
459 let range = interleaved.arrays[array].element_range(row);
460 child_indices.extend(range.map(|i| (array, i)));
461 }
462
463 let child_arrays: Vec<&dyn Array> = interleaved
464 .arrays
465 .iter()
466 .map(|list| list.values().as_ref())
467 .collect();
468 interleave(&child_arrays, &child_indices)
469}
470
471fn interleave_list<O: OffsetSizeTrait>(
472 values: &[&dyn Array],
473 indices: &[(usize, usize)],
474 field: &FieldRef,
475) -> Result<ArrayRef, ArrowError> {
476 let interleaved = Interleave::<'_, GenericListArray<O>>::new(values, indices);
477
478 let mut capacity = 0usize;
480 let mut offsets = Vec::with_capacity(indices.len() + 1);
481 offsets.push(O::from_usize(0).unwrap());
482 for (array, row) in indices {
483 let o = interleaved.arrays[*array].value_offsets();
484 let element_len = o[*row + 1].as_usize() - o[*row].as_usize();
485 capacity += element_len;
486 offsets.push(
487 O::from_usize(capacity).ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?,
488 );
489 }
490
491 macro_rules! list_primitive_helper {
493 ($t:ty) => {
494 interleave_list_like_primitive_child::<GenericListArray<O>, $t>(
495 &interleaved,
496 indices,
497 capacity,
498 field.data_type(),
499 )
500 };
501 }
502
503 let child_values = downcast_primitive! {
504 field.data_type() => (list_primitive_helper),
508 _ => {
509 interleave_list_like_child(&interleaved, indices, capacity)?
510 }
511 };
512
513 let offsets = OffsetBuffer::new(offsets.into());
514 let list_array =
515 GenericListArray::<O>::new(field.clone(), offsets, child_values, interleaved.nulls);
516
517 Ok(Arc::new(list_array))
518}
519
520fn interleave_fixed_size_list(
521 values: &[&dyn Array],
522 indices: &[(usize, usize)],
523 field: &FieldRef,
524 size: i32,
525) -> Result<ArrayRef, ArrowError> {
526 let interleaved = Interleave::<'_, FixedSizeListArray>::new(values, indices);
527 let capacity = indices.len() * size as usize;
528
529 macro_rules! fsl_primitive_helper {
530 ($t:ty) => {
531 interleave_list_like_primitive_child::<FixedSizeListArray, $t>(
532 &interleaved,
533 indices,
534 capacity,
535 field.data_type(),
536 )
537 };
538 }
539
540 let interleaved_values = downcast_primitive! {
541 field.data_type() => (fsl_primitive_helper),
542 _ => {
543 interleave_list_like_child(&interleaved, indices, capacity)?
544 }
545 };
546
547 let array = FixedSizeListArray::new(field.clone(), size, interleaved_values, interleaved.nulls);
548 Ok(Arc::new(array))
549}
550
551fn interleave_map(
552 values: &[&dyn Array],
553 indices: &[(usize, usize)],
554 field: &FieldRef,
555 ordered: bool,
556) -> Result<ArrayRef, ArrowError> {
557 let interleaved = Interleave::<'_, MapArray>::new(values, indices);
558
559 let mut capacity = 0usize;
560 let mut offsets = Vec::with_capacity(indices.len() + 1);
561 offsets.push(0i32);
562 for &(array, row) in indices {
563 let o = interleaved.arrays[array].value_offsets();
564 let element_len = (o[row + 1] - o[row]) as usize;
565 capacity += element_len;
566 offsets
567 .push(i32::try_from(capacity).map_err(|_| ArrowError::OffsetOverflowError(capacity))?);
568 }
569
570 let mut child_indices = Vec::with_capacity(capacity);
571 for &(array, row) in indices {
572 let o = interleaved.arrays[array].value_offsets();
573 let start = o[row] as usize;
574 let end = o[row + 1] as usize;
575 child_indices.extend((start..end).map(|i| (array, i)));
576 }
577
578 let entries_arrays: Vec<&dyn Array> = interleaved
579 .arrays
580 .iter()
581 .map(|m| m.entries() as &dyn Array)
582 .collect();
583 let interleaved_entries = interleave(&entries_arrays, &child_indices)?;
584
585 let offsets = OffsetBuffer::new(offsets.into());
586 let entries = interleaved_entries.as_struct().clone();
587 let array = MapArray::new(field.clone(), offsets, entries, interleaved.nulls, ordered);
588 Ok(Arc::new(array))
589}
590
591fn interleave_run_end<R: RunEndIndexType>(
593 values: &[&dyn Array],
594 indices: &[(usize, usize)],
595) -> Result<ArrayRef, ArrowError> {
596 if indices.is_empty() {
597 return Ok(new_empty_array(values[0].data_type()));
598 }
599
600 let n = indices.len();
601 R::Native::from_usize(n).ok_or_else(|| {
602 ArrowError::ComputeError(format!(
603 "interleave_run_end: output length {n} does not fit run-end type"
604 ))
605 })?;
606
607 let runs: Vec<&RunArray<R>> = values.iter().map(|a| a.as_run::<R>()).collect();
608 let value_arrays: Vec<&dyn Array> = runs.iter().map(|r| r.values().as_ref()).collect();
609
610 let mut phys_pairs: Vec<(usize, usize)> = vec![(0, 0); n];
613 let mut grouped: Vec<(Vec<R::Native>, Vec<usize>)> =
614 (0..runs.len()).map(|_| (Vec::new(), Vec::new())).collect();
615 for (out_pos, &(arr, row)) in indices.iter().enumerate() {
616 let row = R::Native::from_usize(row).ok_or_else(|| {
617 ArrowError::InvalidArgumentError(format!(
618 "interleave_run_end: row index {row} not representable as run-end type {}",
619 R::DATA_TYPE
620 ))
621 })?;
622 grouped[arr].0.push(row);
623 grouped[arr].1.push(out_pos);
624 }
625 for (arr_idx, (logical_rows, out_positions)) in grouped.into_iter().enumerate() {
626 let phys = runs[arr_idx].get_physical_indices(&logical_rows)?;
627 for (p, out_pos) in phys.iter().zip(out_positions.iter()) {
628 phys_pairs[*out_pos] = (arr_idx, *p);
629 }
630 }
631
632 let mut run_ends_buf: Vec<R::Native> = Vec::with_capacity(n);
637 let mut dedup_pairs: Vec<(usize, usize)> = Vec::with_capacity(n);
638 dedup_pairs.push(phys_pairs[0]);
639 for i in 1..n {
640 if phys_pairs[i] != phys_pairs[i - 1] {
641 run_ends_buf.push(R::Native::from_usize(i).unwrap());
642 dedup_pairs.push(phys_pairs[i]);
643 }
644 }
645 run_ends_buf.push(R::Native::from_usize(n).unwrap());
646
647 let taken_values = interleave(&value_arrays, &dedup_pairs)?;
648 let run_ends = PrimitiveArray::<R>::from_iter_values(run_ends_buf);
649
650 Ok(Arc::new(RunArray::<R>::try_new(
651 &run_ends,
652 taken_values.as_ref(),
653 )?))
654}
655
656fn interleave_list_view<O: OffsetSizeTrait>(
657 values: &[&dyn Array],
658 indices: &[(usize, usize)],
659 field: &FieldRef,
660) -> Result<ArrayRef, ArrowError> {
661 let interleaved = Interleave::<'_, GenericListViewArray<O>>::new(values, indices);
662
663 let concat_cost: usize = interleaved.arrays.iter().map(|lv| lv.values().len()).sum();
668 let per_row_cost: usize = indices
669 .iter()
670 .map(|&(a, r)| interleaved.arrays[a].sizes()[r].as_usize())
671 .sum();
672
673 if per_row_cost <= concat_cost {
674 interleave_list_view_copy::<O>(&interleaved, indices, field)
675 } else {
676 interleave_list_view_concat::<O>(&interleaved, indices, field)
677 }
678}
679
680fn interleave_list_view_copy<O: OffsetSizeTrait>(
682 interleaved: &Interleave<'_, GenericListViewArray<O>>,
683 indices: &[(usize, usize)],
684 field: &FieldRef,
685) -> Result<ArrayRef, ArrowError> {
686 let mut capacity = 0usize;
687 let mut offsets = Vec::with_capacity(indices.len());
688 let mut sizes = Vec::with_capacity(indices.len());
689 for &(array_idx, row_idx) in indices {
690 let list = interleaved.arrays[array_idx];
691 let size = list.sizes()[row_idx].as_usize();
692 offsets.push(
693 O::from_usize(capacity).ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?,
694 );
695 sizes.push(O::from_usize(size).ok_or_else(|| ArrowError::OffsetOverflowError(size))?);
696 capacity += size;
697 }
698
699 let child_data: Vec<_> = interleaved
700 .arrays
701 .iter()
702 .map(|list| list.values().to_data())
703 .collect();
704 let child_data_refs: Vec<_> = child_data.iter().collect();
705 let mut mutable_child = MutableArrayData::new(child_data_refs, false, capacity);
706 for &(array_idx, row_idx) in indices {
707 let list = interleaved.arrays[array_idx];
708 let start = list.offsets()[row_idx].as_usize();
709 let size = list.sizes()[row_idx].as_usize();
710 if size > 0 {
711 mutable_child.try_extend(array_idx, start, start + size)?;
712 }
713 }
714
715 Ok(Arc::new(GenericListViewArray::<O>::new(
716 field.clone(),
717 offsets.into(),
718 sizes.into(),
719 make_array(mutable_child.freeze()),
720 interleaved.nulls.clone(),
721 )))
722}
723
724fn interleave_list_view_concat<O: OffsetSizeTrait>(
727 interleaved: &Interleave<'_, GenericListViewArray<O>>,
728 indices: &[(usize, usize)],
729 field: &FieldRef,
730) -> Result<ArrayRef, ArrowError> {
731 let child_arrays: Vec<&dyn Array> = interleaved
732 .arrays
733 .iter()
734 .map(|lv| lv.values().as_ref())
735 .collect();
736 let mut base_offsets = Vec::with_capacity(interleaved.arrays.len());
737 let mut running = 0usize;
738 for lv in &interleaved.arrays {
739 base_offsets.push(running);
740 running += lv.values().len();
741 }
742 let combined_values = concat(&child_arrays)?;
743
744 let mut new_offsets = Vec::with_capacity(indices.len());
745 let mut new_sizes = Vec::with_capacity(indices.len());
746 for &(array_idx, row_idx) in indices {
747 let lv = interleaved.arrays[array_idx];
748 let adjusted = lv.offsets()[row_idx].as_usize() + base_offsets[array_idx];
749 new_offsets.push(
750 O::from_usize(adjusted).ok_or_else(|| ArrowError::OffsetOverflowError(adjusted))?,
751 );
752 new_sizes.push(lv.sizes()[row_idx]);
753 }
754
755 Ok(Arc::new(GenericListViewArray::<O>::new(
756 field.clone(),
757 new_offsets.into(),
758 new_sizes.into(),
759 combined_values,
760 interleaved.nulls.clone(),
761 )))
762}
763
764fn interleave_fallback(
766 values: &[&dyn Array],
767 indices: &[(usize, usize)],
768) -> Result<ArrayRef, ArrowError> {
769 let arrays: Vec<_> = values.iter().map(|x| x.to_data()).collect();
770 let arrays: Vec<_> = arrays.iter().collect();
771 let mut array_data = MutableArrayData::new(arrays, false, indices.len());
772
773 let mut cur_array = indices[0].0;
774 let mut start_row_idx = indices[0].1;
775 let mut end_row_idx = start_row_idx + 1;
776
777 for (array, row) in indices.iter().skip(1).copied() {
778 if array == cur_array && row == end_row_idx {
779 end_row_idx += 1;
781 continue;
782 }
783
784 array_data.try_extend(cur_array, start_row_idx, end_row_idx)?;
786
787 cur_array = array;
789 start_row_idx = row;
790 end_row_idx = start_row_idx + 1;
791 }
792
793 array_data.try_extend(cur_array, start_row_idx, end_row_idx)?;
795 Ok(make_array(array_data.freeze()))
796}
797
798fn interleave_fallback_dictionary<K: ArrowDictionaryKeyType>(
808 dictionaries: &[&DictionaryArray<K>],
809 indices: &[(usize, usize)],
810) -> Result<ArrayRef, ArrowError> {
811 let relative_offsets: Vec<usize> = dictionaries
812 .iter()
813 .scan(0usize, |offset, dict| {
814 let current = *offset;
815 *offset += dict.values().len();
816 Some(current)
817 })
818 .collect();
819 let all_values: Vec<&dyn Array> = dictionaries.iter().map(|d| d.values().as_ref()).collect();
820 let concatenated_values = concat(&all_values)?;
821
822 let any_nulls = dictionaries.iter().any(|d| d.keys().nulls().is_some());
823 let (new_keys, nulls) = if any_nulls {
824 let mut has_nulls = false;
825 let new_keys: Vec<K::Native> = indices
826 .iter()
827 .map(|(array, row)| {
828 let old_keys = dictionaries[*array].keys();
829 if old_keys.is_valid(*row) {
830 let old_key = old_keys.values()[*row].as_usize();
831 K::Native::from_usize(relative_offsets[*array] + old_key)
832 .expect("key overflow should be checked by caller")
833 } else {
834 has_nulls = true;
835 K::Native::ZERO
836 }
837 })
838 .collect();
839
840 let nulls = if has_nulls {
841 let null_buffer = BooleanBuffer::collect_bool(indices.len(), |i| {
842 let (array, row) = indices[i];
843 dictionaries[array].keys().is_valid(row)
844 });
845 Some(NullBuffer::new(null_buffer))
846 } else {
847 None
848 };
849 (new_keys, nulls)
850 } else {
851 let new_keys: Vec<K::Native> = indices
852 .iter()
853 .map(|(array, row)| {
854 let old_key = dictionaries[*array].keys().values()[*row].as_usize();
855 K::Native::from_usize(relative_offsets[*array] + old_key)
856 .expect("key overflow should be checked by caller")
857 })
858 .collect();
859 (new_keys, None)
860 };
861
862 let keys_array = PrimitiveArray::<K>::new(new_keys.into(), nulls);
863 let array = unsafe { DictionaryArray::new_unchecked(keys_array, concatenated_values) };
865 Ok(Arc::new(array))
866}
867
868pub fn interleave_record_batch(
913 record_batches: &[&RecordBatch],
914 indices: &[(usize, usize)],
915) -> Result<RecordBatch, ArrowError> {
916 let schema = record_batches[0].schema();
917 let columns = (0..schema.fields().len())
918 .map(|i| {
919 let column_values: Vec<&dyn Array> = record_batches
920 .iter()
921 .map(|batch| batch.column(i).as_ref())
922 .collect();
923 interleave(&column_values, indices)
924 })
925 .collect::<Result<Vec<_>, _>>()?;
926 RecordBatch::try_new(schema, columns)
927}
928
929#[cfg(test)]
930mod tests {
931 use super::*;
932 use arrow_array::Int32RunArray;
933 use arrow_array::builder::{
934 GenericListBuilder, Int32Builder, PrimitiveBuilder, PrimitiveRunBuilder,
935 };
936 use arrow_array::types::{Decimal128Type, Int8Type, TimestampMicrosecondType};
937 use arrow_buffer::ScalarBuffer;
938 use arrow_schema::{Field, TimeUnit};
939
940 #[test]
941 fn test_primitive() {
942 let a = Int32Array::from_iter_values([1, 2, 3, 4]);
943 let b = Int32Array::from_iter_values([5, 6, 7]);
944 let c = Int32Array::from_iter_values([8, 9, 10]);
945 let values = interleave(&[&a, &b, &c], &[(0, 3), (0, 3), (2, 2), (2, 0), (1, 1)]).unwrap();
946 let v = values.as_primitive::<Int32Type>();
947 assert_eq!(v.values(), &[4, 4, 10, 8, 6]);
948 }
949
950 #[test]
951 fn test_primitive_nulls() {
952 let a = Int32Array::from_iter_values([1, 2, 3, 4]);
953 let b = Int32Array::from_iter([Some(1), Some(4), None]);
954 let values = interleave(&[&a, &b], &[(0, 1), (1, 2), (1, 2), (0, 3), (0, 2)]).unwrap();
955 let v: Vec<_> = values.as_primitive::<Int32Type>().into_iter().collect();
956 assert_eq!(&v, &[Some(2), None, None, Some(4), Some(3)])
957 }
958
959 #[test]
960 fn test_primitive_empty() {
961 let a = Int32Array::from_iter_values([1, 2, 3, 4]);
962 let v = interleave(&[&a], &[]).unwrap();
963 assert!(v.is_empty());
964 assert_eq!(v.data_type(), &DataType::Int32);
965 }
966
967 #[test]
968 fn test_strings() {
969 let a = StringArray::from_iter_values(["a", "b", "c"]);
970 let b = StringArray::from_iter_values(["hello", "world", "foo"]);
971 let values = interleave(&[&a, &b], &[(0, 2), (0, 2), (1, 0), (1, 1), (0, 1)]).unwrap();
972 let v = values.as_string::<i32>();
973 let values: Vec<_> = v.into_iter().collect();
974 assert_eq!(
975 &values,
976 &[
977 Some("c"),
978 Some("c"),
979 Some("hello"),
980 Some("world"),
981 Some("b")
982 ]
983 )
984 }
985
986 #[test]
987 fn test_interleave_dictionary() {
988 let a = DictionaryArray::<Int32Type>::from_iter(["a", "b", "c", "a", "b"]);
989 let b = DictionaryArray::<Int32Type>::from_iter(["a", "c", "a", "c", "a"]);
990
991 let values =
993 interleave(&[&a, &b], &[(0, 2), (0, 2), (0, 2), (1, 0), (1, 1), (0, 1)]).unwrap();
994 let v = values.as_dictionary::<Int32Type>();
995 assert_eq!(v.values().len(), 5);
996
997 let vc = v.downcast_dict::<StringArray>().unwrap();
998 let collected: Vec<_> = vc.into_iter().map(Option::unwrap).collect();
999 assert_eq!(&collected, &["c", "c", "c", "a", "c", "b"]);
1000
1001 let values = interleave(&[&a, &b], &[(0, 2), (0, 2), (1, 1)]).unwrap();
1003 let v = values.as_dictionary::<Int32Type>();
1004 assert_eq!(v.values().len(), 1);
1005
1006 let vc = v.downcast_dict::<StringArray>().unwrap();
1007 let collected: Vec<_> = vc.into_iter().map(Option::unwrap).collect();
1008 assert_eq!(&collected, &["c", "c", "c"]);
1009 }
1010
1011 #[test]
1012 fn test_interleave_dictionary_nulls() {
1013 let input_1_keys = Int32Array::from_iter_values([0, 2, 1, 3]);
1014 let input_1_values = StringArray::from(vec![Some("foo"), None, Some("bar"), Some("fiz")]);
1015 let input_1 = DictionaryArray::new(input_1_keys, Arc::new(input_1_values));
1016 let input_2: DictionaryArray<Int32Type> = vec![None].into_iter().collect();
1017
1018 let expected = vec![Some("fiz"), None, None, Some("foo")];
1019
1020 let values = interleave(
1021 &[&input_1 as _, &input_2 as _],
1022 &[(0, 3), (0, 2), (1, 0), (0, 0)],
1023 )
1024 .unwrap();
1025 let dictionary = values.as_dictionary::<Int32Type>();
1026 let actual: Vec<Option<&str>> = dictionary
1027 .downcast_dict::<StringArray>()
1028 .unwrap()
1029 .into_iter()
1030 .collect();
1031
1032 assert_eq!(actual, expected);
1033 }
1034
1035 #[test]
1036 fn test_interleave_dictionary_overflow_same_values() {
1037 let values: ArrayRef = Arc::new(StringArray::from_iter_values(
1038 (0..50).map(|i| format!("v{i}")),
1039 ));
1040
1041 let dict1 = DictionaryArray::<Int8Type>::new(
1047 Int8Array::from_iter_values([0, 1, 2]),
1048 values.clone(),
1049 );
1050 let dict2 = DictionaryArray::<Int8Type>::new(
1051 Int8Array::from_iter_values([0, 1, 2]),
1052 values.clone(),
1053 );
1054 let dict3 =
1055 DictionaryArray::<Int8Type>::new(Int8Array::from_iter_values([49]), values.clone());
1056
1057 let indices = &[(0, 0), (1, 0), (2, 0)];
1058 let result = interleave(&[&dict1, &dict2, &dict3], indices).unwrap();
1059
1060 let dict_result = result.as_dictionary::<Int8Type>();
1061 let string_result: Vec<_> = dict_result
1062 .downcast_dict::<StringArray>()
1063 .unwrap()
1064 .into_iter()
1065 .map(|x| x.unwrap())
1066 .collect();
1067 assert_eq!(string_result, vec!["v0", "v0", "v49"]);
1068 }
1069
1070 fn test_interleave_lists<O: OffsetSizeTrait>() {
1071 let mut a = GenericListBuilder::<O, _>::new(Int32Builder::new());
1073 a.values().append_value(1);
1074 a.values().append_value(2);
1075 a.append(true);
1076 a.append(false);
1077 a.values().append_value(3);
1078 a.append(true);
1079 let a = a.finish();
1080
1081 let mut b = GenericListBuilder::<O, _>::new(Int32Builder::new());
1083 b.values().append_value(4);
1084 b.append(true);
1085 b.append(false);
1086 b.values().append_value(5);
1087 b.values().append_value(6);
1088 b.values().append_null();
1089 b.append(true);
1090 let b = b.finish();
1091
1092 let values = interleave(&[&a, &b], &[(0, 2), (0, 1), (1, 0), (1, 2), (1, 1)]).unwrap();
1093 let v = values
1094 .as_any()
1095 .downcast_ref::<GenericListArray<O>>()
1096 .unwrap();
1097
1098 let mut expected = GenericListBuilder::<O, _>::new(Int32Builder::new());
1100 expected.values().append_value(3);
1101 expected.append(true);
1102 expected.append(false);
1103 expected.values().append_value(4);
1104 expected.append(true);
1105 expected.values().append_value(5);
1106 expected.values().append_value(6);
1107 expected.values().append_null();
1108 expected.append(true);
1109 expected.append(false);
1110 let expected = expected.finish();
1111
1112 assert_eq!(v, &expected);
1113 }
1114
1115 #[test]
1116 fn test_lists() {
1117 test_interleave_lists::<i32>();
1118 }
1119
1120 #[test]
1121 fn test_large_lists() {
1122 test_interleave_lists::<i64>();
1123 }
1124
1125 type ListRow<T> = Option<Vec<Option<<T as ArrowPrimitiveType>::Native>>>;
1128
1129 fn list_of_primitive<O: OffsetSizeTrait, T: ArrowPrimitiveType>(
1132 data_type: &DataType,
1133 rows: &[ListRow<T>],
1134 ) -> GenericListArray<O> {
1135 let mut builder = GenericListBuilder::<O, _>::new(
1136 PrimitiveBuilder::<T>::new().with_data_type(data_type.clone()),
1137 );
1138 for row in rows {
1139 match row {
1140 Some(items) => {
1141 items
1142 .iter()
1143 .for_each(|v| builder.values().append_option(*v));
1144 builder.append(true);
1145 }
1146 None => builder.append(false),
1147 }
1148 }
1149 builder.finish()
1150 }
1151
1152 fn check_interleave_list_primitive<O: OffsetSizeTrait, T: ArrowPrimitiveType>(
1155 data_type: &DataType,
1156 inputs: &[&[ListRow<T>]],
1157 indices: &[(usize, usize)],
1158 expected: &[ListRow<T>],
1159 ) {
1160 let arrays: Vec<_> = inputs
1161 .iter()
1162 .map(|rows| list_of_primitive::<O, T>(data_type, rows))
1163 .collect();
1164 let refs: Vec<&dyn Array> = arrays.iter().map(|a| a as &dyn Array).collect();
1165
1166 let values = interleave(&refs, indices).unwrap();
1167 let v = values
1168 .as_any()
1169 .downcast_ref::<GenericListArray<O>>()
1170 .unwrap();
1171
1172 assert_eq!(v, &list_of_primitive::<O, T>(data_type, expected));
1173 assert_eq!(v.values().data_type(), data_type);
1176 }
1177
1178 fn test_interleave_lists_decimal<O: OffsetSizeTrait>() {
1179 check_interleave_list_primitive::<O, Decimal128Type>(
1181 &DataType::Decimal128(20, 3),
1182 &[
1183 &[
1184 Some(vec![Some(1), Some(2)]),
1185 None,
1186 Some(vec![Some(3), None]),
1187 ], &[Some(vec![Some(4)]), Some(vec![Some(5), Some(6)])], ],
1190 &[(0, 2), (0, 1), (1, 0), (1, 1)],
1191 &[
1192 Some(vec![Some(3), None]),
1193 None,
1194 Some(vec![Some(4)]),
1195 Some(vec![Some(5), Some(6)]),
1196 ],
1197 );
1198 }
1199
1200 #[test]
1201 fn test_lists_decimal() {
1202 test_interleave_lists_decimal::<i32>();
1203 test_interleave_lists_decimal::<i64>();
1204 }
1205
1206 fn test_interleave_lists_timestamp_tz<O: OffsetSizeTrait>() {
1207 check_interleave_list_primitive::<O, TimestampMicrosecondType>(
1209 &DataType::Timestamp(TimeUnit::Microsecond, Some("+08:00".into())),
1210 &[&[Some(vec![Some(1), Some(2)]), Some(vec![Some(3)])]],
1211 &[(0, 1), (0, 0)],
1212 &[Some(vec![Some(3)]), Some(vec![Some(1), Some(2)])],
1213 );
1214 }
1215
1216 #[test]
1217 fn test_lists_timestamp_tz() {
1218 test_interleave_lists_timestamp_tz::<i32>();
1219 test_interleave_lists_timestamp_tz::<i64>();
1220 }
1221
1222 fn test_interleave_list_views<O: OffsetSizeTrait>() {
1223 let mut a = GenericListBuilder::<O, _>::new(Int32Builder::new());
1225 a.values().append_value(1);
1226 a.values().append_value(2);
1227 a.append(true);
1228 a.append(false);
1229 a.values().append_value(3);
1230 a.append(true);
1231 let a: GenericListViewArray<O> = a.finish().into();
1232
1233 let mut b = GenericListBuilder::<O, _>::new(Int32Builder::new());
1235 b.values().append_value(4);
1236 b.append(true);
1237 b.append(false);
1238 b.values().append_value(5);
1239 b.values().append_value(6);
1240 b.values().append_null();
1241 b.append(true);
1242 let b: GenericListViewArray<O> = b.finish().into();
1243
1244 let values = interleave(&[&a, &b], &[(0, 2), (0, 1), (1, 0), (1, 2), (1, 1)]).unwrap();
1245 let v = values
1246 .as_any()
1247 .downcast_ref::<GenericListViewArray<O>>()
1248 .unwrap();
1249
1250 let mut expected = GenericListBuilder::<O, _>::new(Int32Builder::new());
1252 expected.values().append_value(3);
1253 expected.append(true);
1254 expected.append(false);
1255 expected.values().append_value(4);
1256 expected.append(true);
1257 expected.values().append_value(5);
1258 expected.values().append_value(6);
1259 expected.values().append_null();
1260 expected.append(true);
1261 expected.append(false);
1262 let expected: GenericListViewArray<O> = expected.finish().into();
1263
1264 assert_eq!(v, &expected);
1265 }
1266
1267 #[test]
1268 fn test_list_views() {
1269 test_interleave_list_views::<i32>();
1270 }
1271
1272 #[test]
1273 fn test_large_list_views() {
1274 test_interleave_list_views::<i64>();
1275 }
1276
1277 #[test]
1278 fn test_interleave_list_view_overlapping() {
1279 let field = Arc::new(Field::new_list_field(DataType::Int64, false));
1280
1281 let lv_a = ListViewArray::new(
1285 Arc::clone(&field),
1286 ScalarBuffer::from(vec![0i32, 0, 0, 0, 0, 5, 5, 5, 5, 5]),
1287 ScalarBuffer::from(vec![5i32; 10]),
1288 Arc::new(Int64Array::from_iter_values(0..10)),
1289 None,
1290 );
1291
1292 let lv_b = ListViewArray::new(
1296 Arc::clone(&field),
1297 ScalarBuffer::from(vec![0i32, 0, 0, 0, 3, 3, 3, 3]),
1298 ScalarBuffer::from(vec![3i32; 8]),
1299 Arc::new(Int64Array::from_iter_values(100..106)),
1300 None,
1301 );
1302
1303 let indices: Vec<(usize, usize)> = vec![
1304 (0, 0),
1305 (1, 0),
1306 (0, 5),
1307 (1, 4),
1308 (0, 1),
1309 (1, 1),
1310 (0, 6),
1311 (1, 5),
1312 ];
1313 let result = interleave(&[&lv_a as &dyn Array, &lv_b as &dyn Array], &indices).unwrap();
1314 result
1315 .to_data()
1316 .validate_full()
1317 .expect("result must be valid");
1318
1319 let result_lv = result.as_list_view::<i32>();
1320 assert_eq!(result_lv.len(), 8);
1321 assert_eq!(
1322 result_lv.value(0).as_primitive::<Int64Type>().values(),
1323 &[0, 1, 2, 3, 4]
1324 );
1325 assert_eq!(
1326 result_lv.value(1).as_primitive::<Int64Type>().values(),
1327 &[100, 101, 102]
1328 );
1329 assert_eq!(
1330 result_lv.value(2).as_primitive::<Int64Type>().values(),
1331 &[5, 6, 7, 8, 9]
1332 );
1333 assert_eq!(
1334 result_lv.value(3).as_primitive::<Int64Type>().values(),
1335 &[103, 104, 105]
1336 );
1337
1338 let total_input_elements = lv_a.values().len() + lv_b.values().len();
1341 assert_eq!(result_lv.values().len(), total_input_elements);
1342 }
1343
1344 #[test]
1345 fn test_struct_without_nulls() {
1346 let fields = Fields::from(vec![
1347 Field::new("number_col", DataType::Int32, false),
1348 Field::new("string_col", DataType::Utf8, false),
1349 ]);
1350 let a = {
1351 let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
1352 let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);
1353
1354 StructArray::try_new(
1355 fields.clone(),
1356 vec![Arc::new(number_col), Arc::new(string_col)],
1357 None,
1358 )
1359 .unwrap()
1360 };
1361
1362 let b = {
1363 let number_col = Int32Array::from_iter_values([5, 6, 7]);
1364 let string_col = StringArray::from_iter_values(["hello", "world", "foo"]);
1365
1366 StructArray::try_new(
1367 fields.clone(),
1368 vec![Arc::new(number_col), Arc::new(string_col)],
1369 None,
1370 )
1371 .unwrap()
1372 };
1373
1374 let c = {
1375 let number_col = Int32Array::from_iter_values([8, 9, 10]);
1376 let string_col = StringArray::from_iter_values(["x", "y", "z"]);
1377
1378 StructArray::try_new(
1379 fields.clone(),
1380 vec![Arc::new(number_col), Arc::new(string_col)],
1381 None,
1382 )
1383 .unwrap()
1384 };
1385
1386 let values = interleave(&[&a, &b, &c], &[(0, 3), (0, 3), (2, 2), (2, 0), (1, 1)]).unwrap();
1387 let values_struct = values.as_struct();
1388 assert_eq!(values_struct.data_type(), &DataType::Struct(fields));
1389 assert_eq!(values_struct.null_count(), 0);
1390
1391 let values_number = values_struct.column(0).as_primitive::<Int32Type>();
1392 assert_eq!(values_number.values(), &[4, 4, 10, 8, 6]);
1393 let values_string = values_struct.column(1).as_string::<i32>();
1394 let values_string: Vec<_> = values_string.into_iter().collect();
1395 assert_eq!(
1396 &values_string,
1397 &[Some("d"), Some("d"), Some("z"), Some("x"), Some("world")]
1398 );
1399 }
1400
1401 #[test]
1402 fn test_struct_with_nulls_in_values() {
1403 let fields = Fields::from(vec![
1404 Field::new("number_col", DataType::Int32, true),
1405 Field::new("string_col", DataType::Utf8, true),
1406 ]);
1407 let a = {
1408 let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
1409 let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);
1410
1411 StructArray::try_new(
1412 fields.clone(),
1413 vec![Arc::new(number_col), Arc::new(string_col)],
1414 None,
1415 )
1416 .unwrap()
1417 };
1418
1419 let b = {
1420 let number_col = Int32Array::from_iter([Some(1), Some(4), None]);
1421 let string_col = StringArray::from(vec![Some("hello"), None, Some("foo")]);
1422
1423 StructArray::try_new(
1424 fields.clone(),
1425 vec![Arc::new(number_col), Arc::new(string_col)],
1426 None,
1427 )
1428 .unwrap()
1429 };
1430
1431 let values = interleave(&[&a, &b], &[(0, 1), (1, 2), (1, 2), (0, 3), (1, 1)]).unwrap();
1432 let values_struct = values.as_struct();
1433 assert_eq!(values_struct.data_type(), &DataType::Struct(fields));
1434
1435 assert_eq!(values_struct.null_count(), 0);
1437
1438 let values_number: Vec<_> = values_struct
1439 .column(0)
1440 .as_primitive::<Int32Type>()
1441 .into_iter()
1442 .collect();
1443 assert_eq!(values_number, &[Some(2), None, None, Some(4), Some(4)]);
1444
1445 let values_string = values_struct.column(1).as_string::<i32>();
1446 let values_string: Vec<_> = values_string.into_iter().collect();
1447 assert_eq!(
1448 &values_string,
1449 &[Some("b"), Some("foo"), Some("foo"), Some("d"), None]
1450 );
1451 }
1452
1453 #[test]
1454 fn test_struct_with_nulls() {
1455 let fields = Fields::from(vec![
1456 Field::new("number_col", DataType::Int32, false),
1457 Field::new("string_col", DataType::Utf8, false),
1458 ]);
1459 let a = {
1460 let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
1461 let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);
1462
1463 StructArray::try_new(
1464 fields.clone(),
1465 vec![Arc::new(number_col), Arc::new(string_col)],
1466 None,
1467 )
1468 .unwrap()
1469 };
1470
1471 let b = {
1472 let number_col = Int32Array::from_iter_values([5, 6, 7]);
1473 let string_col = StringArray::from_iter_values(["hello", "world", "foo"]);
1474
1475 StructArray::try_new(
1476 fields.clone(),
1477 vec![Arc::new(number_col), Arc::new(string_col)],
1478 Some(NullBuffer::from(&[true, false, true])),
1479 )
1480 .unwrap()
1481 };
1482
1483 let c = {
1484 let number_col = Int32Array::from_iter_values([8, 9, 10]);
1485 let string_col = StringArray::from_iter_values(["x", "y", "z"]);
1486
1487 StructArray::try_new(
1488 fields.clone(),
1489 vec![Arc::new(number_col), Arc::new(string_col)],
1490 None,
1491 )
1492 .unwrap()
1493 };
1494
1495 let values = interleave(&[&a, &b, &c], &[(0, 3), (0, 3), (2, 2), (1, 1), (2, 0)]).unwrap();
1496 let values_struct = values.as_struct();
1497 assert_eq!(values_struct.data_type(), &DataType::Struct(fields));
1498
1499 let validity: Vec<bool> = {
1500 let null_buffer = values_struct.nulls().expect("should_have_nulls");
1501
1502 null_buffer.iter().collect()
1503 };
1504 assert_eq!(validity, &[true, true, true, false, true]);
1505 let values_number = values_struct.column(0).as_primitive::<Int32Type>();
1506 assert_eq!(values_number.values(), &[4, 4, 10, 6, 8]);
1507 let values_string = values_struct.column(1).as_string::<i32>();
1508 let values_string: Vec<_> = values_string.into_iter().collect();
1509 assert_eq!(
1510 &values_string,
1511 &[Some("d"), Some("d"), Some("z"), Some("world"), Some("x"),]
1512 );
1513 }
1514
1515 #[test]
1516 fn test_struct_empty() {
1517 let fields = Fields::from(vec![
1518 Field::new("number_col", DataType::Int32, false),
1519 Field::new("string_col", DataType::Utf8, false),
1520 ]);
1521 let a = {
1522 let number_col = Int32Array::from_iter_values([1, 2, 3, 4]);
1523 let string_col = StringArray::from_iter_values(["a", "b", "c", "d"]);
1524
1525 StructArray::try_new(
1526 fields.clone(),
1527 vec![Arc::new(number_col), Arc::new(string_col)],
1528 None,
1529 )
1530 .unwrap()
1531 };
1532 let v = interleave(&[&a], &[]).unwrap();
1533 assert!(v.is_empty());
1534 assert_eq!(v.data_type(), &DataType::Struct(fields));
1535 }
1536
1537 #[test]
1538 fn interleave_sparse_nulls() {
1539 let values = StringArray::from_iter_values((0..100).map(|x| x.to_string()));
1540 let keys = Int32Array::from_iter_values(0..10);
1541 let dict_a = DictionaryArray::new(keys, Arc::new(values));
1542 let values = StringArray::new_null(0);
1543 let keys = Int32Array::new_null(10);
1544 let dict_b = DictionaryArray::new(keys, Arc::new(values));
1545
1546 let indices = &[(0, 0), (0, 1), (0, 2), (1, 0)];
1547 let array = interleave(&[&dict_a, &dict_b], indices).unwrap();
1548
1549 let expected =
1550 DictionaryArray::<Int32Type>::from_iter(vec![Some("0"), Some("1"), Some("2"), None]);
1551 assert_eq!(array.as_ref(), &expected)
1552 }
1553
1554 #[test]
1555 fn test_interleave_views() {
1556 let values = StringArray::from_iter_values([
1557 "hello",
1558 "world_long_string_not_inlined",
1559 "foo",
1560 "bar",
1561 "baz",
1562 ]);
1563 let view_a = StringViewArray::from(&values);
1564
1565 let values = StringArray::from_iter_values([
1566 "test",
1567 "data",
1568 "more_long_string_not_inlined",
1569 "views",
1570 "here",
1571 ]);
1572 let view_b = StringViewArray::from(&values);
1573
1574 let indices = &[
1575 (0, 2), (1, 0), (0, 4), (1, 3), (0, 1), ];
1581
1582 let values = interleave(&[&view_a, &view_b], indices).unwrap();
1584 let result = values.as_string_view();
1585 assert_eq!(result.data_buffers().len(), 1);
1586
1587 let fallback = interleave_fallback(&[&view_a, &view_b], indices).unwrap();
1588 let fallback_result = fallback.as_string_view();
1589 assert_eq!(fallback_result.data_buffers().len(), 2);
1591
1592 let collected: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();
1594
1595 let fallback_collected: Vec<_> = fallback_result
1596 .iter()
1597 .map(|x| x.map(|s| s.to_string()))
1598 .collect();
1599
1600 assert_eq!(&collected, &fallback_collected);
1601
1602 assert_eq!(
1603 &collected,
1604 &[
1605 Some("foo".to_string()),
1606 Some("test".to_string()),
1607 Some("baz".to_string()),
1608 Some("views".to_string()),
1609 Some("world_long_string_not_inlined".to_string()),
1610 ]
1611 );
1612 }
1613
1614 #[test]
1615 fn test_interleave_views_with_nulls() {
1616 let values = StringArray::from_iter([
1617 Some("hello"),
1618 None,
1619 Some("foo_long_string_not_inlined"),
1620 Some("bar"),
1621 None,
1622 ]);
1623 let view_a = StringViewArray::from(&values);
1624
1625 let values = StringArray::from_iter([
1626 Some("test"),
1627 Some("data_long_string_not_inlined"),
1628 None,
1629 None,
1630 Some("here"),
1631 ]);
1632 let view_b = StringViewArray::from(&values);
1633
1634 let indices = &[
1635 (0, 1), (1, 2), (0, 2), (1, 3), (0, 4), ];
1641
1642 let values = interleave(&[&view_a, &view_b], indices).unwrap();
1644 let result = values.as_string_view();
1645 assert_eq!(result.data_buffers().len(), 1);
1646
1647 let fallback = interleave_fallback(&[&view_a, &view_b], indices).unwrap();
1648 let fallback_result = fallback.as_string_view();
1649
1650 let collected: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();
1652
1653 let fallback_collected: Vec<_> = fallback_result
1654 .iter()
1655 .map(|x| x.map(|s| s.to_string()))
1656 .collect();
1657
1658 assert_eq!(&collected, &fallback_collected);
1659
1660 assert_eq!(
1661 &collected,
1662 &[
1663 None,
1664 None,
1665 Some("foo_long_string_not_inlined".to_string()),
1666 None,
1667 None,
1668 ]
1669 );
1670 }
1671
1672 #[test]
1673 fn test_interleave_views_multiple_buffers() {
1674 let str1 = "very_long_string_from_first_buffer".as_bytes();
1675 let str2 = "very_long_string_from_second_buffer".as_bytes();
1676 let buffer1 = str1.to_vec().into();
1677 let buffer2 = str2.to_vec().into();
1678
1679 let view1 = ByteView::new(str1.len() as u32, &str1[..4])
1680 .with_buffer_index(0)
1681 .with_offset(0)
1682 .as_u128();
1683 let view2 = ByteView::new(str2.len() as u32, &str2[..4])
1684 .with_buffer_index(1)
1685 .with_offset(0)
1686 .as_u128();
1687 let view_a =
1688 StringViewArray::try_new(vec![view1, view2].into(), vec![buffer1, buffer2], None)
1689 .unwrap();
1690
1691 let str3 = "another_very_long_string_buffer_three".as_bytes();
1692 let str4 = "different_long_string_in_buffer_four".as_bytes();
1693 let buffer3 = str3.to_vec().into();
1694 let buffer4 = str4.to_vec().into();
1695
1696 let view3 = ByteView::new(str3.len() as u32, &str3[..4])
1697 .with_buffer_index(0)
1698 .with_offset(0)
1699 .as_u128();
1700 let view4 = ByteView::new(str4.len() as u32, &str4[..4])
1701 .with_buffer_index(1)
1702 .with_offset(0)
1703 .as_u128();
1704 let view_b =
1705 StringViewArray::try_new(vec![view3, view4].into(), vec![buffer3, buffer4], None)
1706 .unwrap();
1707
1708 let indices = &[
1709 (0, 0), (1, 0), (0, 1), (1, 1), (0, 0), (1, 1), ];
1716
1717 let values = interleave(&[&view_a, &view_b], indices).unwrap();
1719 let result = values.as_string_view();
1720
1721 assert_eq!(
1722 result.data_buffers().len(),
1723 4,
1724 "Expected four buffers (two from each input array)"
1725 );
1726
1727 let result_strings: Vec<_> = result.iter().map(|x| x.map(|s| s.to_string())).collect();
1728 assert_eq!(
1729 result_strings,
1730 vec![
1731 Some("very_long_string_from_first_buffer".to_string()),
1732 Some("another_very_long_string_buffer_three".to_string()),
1733 Some("very_long_string_from_second_buffer".to_string()),
1734 Some("different_long_string_in_buffer_four".to_string()),
1735 Some("very_long_string_from_first_buffer".to_string()),
1736 Some("different_long_string_in_buffer_four".to_string()),
1737 ]
1738 );
1739
1740 let views = result.views();
1741 let buffer_indices: Vec<_> = views
1742 .iter()
1743 .map(|raw_view| ByteView::from(*raw_view).buffer_index)
1744 .collect();
1745
1746 assert_eq!(
1747 buffer_indices,
1748 vec![
1749 0, 1, 2, 3, 0, 3, ]
1756 );
1757 }
1758
1759 #[test]
1760 fn test_interleave_run_end_encoded_primitive() {
1761 let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
1762 builder.extend([1, 1, 2, 2, 2, 3].into_iter().map(Some));
1763 let a = builder.finish();
1764
1765 let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
1766 builder.extend([4, 5, 5, 6, 6, 6].into_iter().map(Some));
1767 let b = builder.finish();
1768
1769 let indices = &[(0, 1), (1, 0), (0, 4), (1, 2), (0, 5)];
1770 let result = interleave(&[&a, &b], indices).unwrap();
1771
1772 assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
1774
1775 let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();
1777
1778 let expected = vec![1, 4, 2, 5, 3];
1780 let mut actual = Vec::new();
1781 for i in 0..result_run_array.len() {
1782 let physical_idx = result_run_array.get_physical_index(i);
1783 let value = result_run_array
1784 .values()
1785 .as_primitive::<Int32Type>()
1786 .value(physical_idx);
1787 actual.push(value);
1788 }
1789 assert_eq!(actual, expected);
1790 }
1791
1792 #[test]
1793 fn test_interleave_run_end_encoded_sliced() {
1794 let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
1795 builder.extend([1, 1, 2, 2, 2, 3].into_iter().map(Some));
1796 let a = builder.finish();
1797 let a = a.slice(2, 3); let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
1800 builder.extend([4, 5, 5, 6, 6, 6].into_iter().map(Some));
1801 let b = builder.finish();
1802 let b = b.slice(1, 3); let indices = &[(0, 1), (1, 0), (0, 2), (1, 1), (1, 2)];
1805 let result = interleave(&[&a, &b], indices).unwrap();
1806
1807 let result = result.as_run::<Int32Type>();
1808 let result = result.downcast::<Int32Array>().unwrap();
1809
1810 let expected = vec![2, 5, 2, 5, 6];
1811 let actual = result.into_iter().flatten().collect::<Vec<_>>();
1812 assert_eq!(actual, expected);
1813 }
1814
1815 #[test]
1816 fn test_interleave_run_end_encoded_string() {
1817 let a: Int32RunArray = vec!["hello", "hello", "world", "world", "foo"]
1818 .into_iter()
1819 .collect();
1820 let b: Int32RunArray = vec!["bar", "baz", "baz", "qux"].into_iter().collect();
1821
1822 let indices = &[(0, 0), (1, 1), (0, 3), (1, 3), (0, 4)];
1823 let result = interleave(&[&a, &b], indices).unwrap();
1824
1825 assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
1827
1828 let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();
1830
1831 let expected = vec!["hello", "baz", "world", "qux", "foo"];
1833 let mut actual = Vec::new();
1834 for i in 0..result_run_array.len() {
1835 let physical_idx = result_run_array.get_physical_index(i);
1836 let value = result_run_array
1837 .values()
1838 .as_string::<i32>()
1839 .value(physical_idx);
1840 actual.push(value);
1841 }
1842 assert_eq!(actual, expected);
1843 }
1844
1845 #[test]
1846 fn test_interleave_run_end_encoded_with_nulls() {
1847 let a: Int32RunArray = vec![Some("a"), Some("a"), None, None, Some("b")]
1848 .into_iter()
1849 .collect();
1850 let b: Int32RunArray = vec![None, Some("c"), Some("c"), Some("d")]
1851 .into_iter()
1852 .collect();
1853
1854 let indices = &[(0, 1), (1, 0), (0, 2), (1, 3), (0, 4)];
1855 let result = interleave(&[&a, &b], indices).unwrap();
1856
1857 assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
1859
1860 let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();
1862
1863 let expected = vec![Some("a"), None, None, Some("d"), Some("b")];
1865 let mut actual = Vec::new();
1866 for i in 0..result_run_array.len() {
1867 let physical_idx = result_run_array.get_physical_index(i);
1868 if result_run_array.values().is_null(physical_idx) {
1869 actual.push(None);
1870 } else {
1871 let value = result_run_array
1872 .values()
1873 .as_string::<i32>()
1874 .value(physical_idx);
1875 actual.push(Some(value));
1876 }
1877 }
1878 assert_eq!(actual, expected);
1879 }
1880
1881 #[test]
1882 fn test_interleave_run_end_encoded_different_run_types() {
1883 let mut builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
1884 builder.extend([1, 1, 2, 3, 3].into_iter().map(Some));
1885 let a = builder.finish();
1886
1887 let mut builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
1888 builder.extend([4, 5, 5, 6].into_iter().map(Some));
1889 let b = builder.finish();
1890
1891 let indices = &[(0, 0), (1, 1), (0, 3), (1, 3)];
1892 let result = interleave(&[&a, &b], indices).unwrap();
1893
1894 assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
1896
1897 let result_run_array: &RunArray<Int16Type> = result.as_any().downcast_ref().unwrap();
1899
1900 let expected = vec![1, 5, 3, 6];
1902 let mut actual = Vec::new();
1903 for i in 0..result_run_array.len() {
1904 let physical_idx = result_run_array.get_physical_index(i);
1905 let value = result_run_array
1906 .values()
1907 .as_primitive::<Int32Type>()
1908 .value(physical_idx);
1909 actual.push(value);
1910 }
1911 assert_eq!(actual, expected);
1912 }
1913
1914 #[test]
1915 fn test_interleave_run_end_encoded_mixed_run_lengths() {
1916 let mut builder = PrimitiveRunBuilder::<Int64Type, Int32Type>::new();
1917 builder.extend([1, 2, 2, 2, 2, 3, 3, 4].into_iter().map(Some));
1918 let a = builder.finish();
1919
1920 let mut builder = PrimitiveRunBuilder::<Int64Type, Int32Type>::new();
1921 builder.extend([5, 5, 5, 6, 7, 7, 8, 8].into_iter().map(Some));
1922 let b = builder.finish();
1923
1924 let indices = &[
1925 (0, 0), (1, 2), (0, 3), (1, 3), (0, 6), (1, 6), (0, 7), (1, 4), ];
1934 let result = interleave(&[&a, &b], indices).unwrap();
1935
1936 assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
1938
1939 let result_run_array: &RunArray<Int64Type> = result.as_any().downcast_ref().unwrap();
1941
1942 let expected = vec![1, 5, 2, 6, 3, 8, 4, 7];
1944 let mut actual = Vec::new();
1945 for i in 0..result_run_array.len() {
1946 let physical_idx = result_run_array.get_physical_index(i);
1947 let value = result_run_array
1948 .values()
1949 .as_primitive::<Int32Type>()
1950 .value(physical_idx);
1951 actual.push(value);
1952 }
1953 assert_eq!(actual, expected);
1954 }
1955
1956 #[test]
1957 fn test_interleave_run_end_encoded_empty_runs() {
1958 let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
1959 builder.extend([1].into_iter().map(Some));
1960 let a = builder.finish();
1961
1962 let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
1963 builder.extend([2, 2, 2].into_iter().map(Some));
1964 let b = builder.finish();
1965
1966 let indices = &[(0, 0), (1, 1), (1, 2)];
1967 let result = interleave(&[&a, &b], indices).unwrap();
1968
1969 assert!(matches!(result.data_type(), DataType::RunEndEncoded(_, _)));
1971
1972 let result_run_array: &Int32RunArray = result.as_any().downcast_ref().unwrap();
1974
1975 let expected = vec![1, 2, 2];
1977 let mut actual = Vec::new();
1978 for i in 0..result_run_array.len() {
1979 let physical_idx = result_run_array.get_physical_index(i);
1980 let value = result_run_array
1981 .values()
1982 .as_primitive::<Int32Type>()
1983 .value(physical_idx);
1984 actual.push(value);
1985 }
1986 assert_eq!(actual, expected);
1987 }
1988
1989 #[test]
1990 fn test_struct_no_fields() {
1991 let fields = Fields::empty();
1992 let a = StructArray::try_new_with_length(fields.clone(), vec![], None, 10).unwrap();
1993 let v = interleave(&[&a], &[(0, 0)]).unwrap();
1994 assert_eq!(v.len(), 1);
1995 assert_eq!(v.data_type(), &DataType::Struct(fields));
1996 }
1997
1998 #[test]
1999 fn test_interleave_fallback_dictionary_with_nulls() {
2000 let input_1_keys = Int32Array::from_iter([Some(0), None, Some(1)]);
2001 let input_1_values = StringArray::from_iter_values(["foo", "bar"]);
2002 let dict_a = DictionaryArray::new(input_1_keys, Arc::new(input_1_values));
2003
2004 let input_2_keys = Int32Array::from_iter([Some(0), Some(1), None]);
2005 let input_2_values = StringArray::from_iter_values(["baz", "qux"]);
2006 let dict_b = DictionaryArray::new(input_2_keys, Arc::new(input_2_values));
2007
2008 let indices = vec![
2009 (0, 0), (0, 1), (1, 0), (1, 2), (0, 2), (1, 1), ];
2016
2017 let result =
2018 interleave_fallback_dictionary::<Int32Type>(&[&dict_a, &dict_b], &indices).unwrap();
2019 let dict_result = result.as_dictionary::<Int32Type>();
2020
2021 let string_result = dict_result.downcast_dict::<StringArray>().unwrap();
2022 let collected: Vec<_> = string_result.into_iter().collect();
2023 assert_eq!(
2024 collected,
2025 vec![
2026 Some("foo"),
2027 None,
2028 Some("baz"),
2029 None,
2030 Some("bar"),
2031 Some("qux")
2032 ]
2033 );
2034 }
2035
2036 #[test]
2037 fn test_interleave_bytes_offset_overflow() {
2038 let indices: Vec<(usize, usize)> = vec![(0, 0); (i32::MAX >> 4) as usize];
2039 let text = ('a'..='z').collect::<String>();
2040 let values = StringArray::from(vec![Some(text)]);
2041 assert!(matches!(
2042 interleave(&[&values], &indices),
2043 Err(ArrowError::OffsetOverflowError(_))
2044 ));
2045 }
2046
2047 #[test]
2048 fn test_interleave_list_offset_overflow() {
2049 let mut builder = GenericListBuilder::<i32, _>::new(Int32Builder::new());
2051 for i in 0..32 {
2052 builder.values().append_value(i);
2053 }
2054 builder.append(true);
2055 let list = builder.finish();
2056
2057 let indices: Vec<(usize, usize)> = vec![(0, 0); (i32::MAX as usize / 32) + 1];
2059 assert!(matches!(
2060 interleave(&[&list], &indices),
2061 Err(ArrowError::OffsetOverflowError(_))
2062 ));
2063 }
2064
2065 #[test]
2066 fn test_interleave_list_view() {
2067 let field = Arc::new(Field::new_list_field(DataType::Int64, false));
2076
2077 let lv_a = ListViewArray::new(
2078 Arc::clone(&field),
2079 ScalarBuffer::from(vec![0i32, 2]),
2080 ScalarBuffer::from(vec![2i32, 1]),
2081 Arc::new(Int64Array::from(vec![1_i64, 2, 3])),
2082 None,
2083 );
2084 let lv_b = ListViewArray::new(
2085 field,
2086 ScalarBuffer::from(vec![0i32]),
2087 ScalarBuffer::from(vec![3i32]),
2088 Arc::new(Int64Array::from(vec![4_i64, 5, 6])),
2089 None,
2090 );
2091
2092 let result = interleave(
2093 &[&lv_a as &dyn Array, &lv_b as &dyn Array],
2094 &[(0, 0), (1, 0), (0, 1)],
2095 )
2096 .unwrap();
2097
2098 result
2099 .to_data()
2100 .validate_full()
2101 .expect("interleaved ListViewArray must be internally consistent");
2102
2103 let result_lv = result.as_list_view::<i32>();
2104 assert_eq!(result_lv.len(), 3);
2105 assert_eq!(
2106 result_lv.value(0).as_primitive::<Int64Type>().values(),
2107 &[1, 2]
2108 );
2109 assert_eq!(
2110 result_lv.value(1).as_primitive::<Int64Type>().values(),
2111 &[4, 5, 6]
2112 );
2113 assert_eq!(
2114 result_lv.value(2).as_primitive::<Int64Type>().values(),
2115 &[3]
2116 );
2117 }
2118
2119 #[test]
2120 fn test_interleave_fixed_size_list() {
2121 let field = Arc::new(Field::new("item", DataType::Int32, false));
2123 let a = FixedSizeListArray::new(
2124 field.clone(),
2125 2,
2126 Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])),
2127 None,
2128 );
2129 let b = FixedSizeListArray::new(
2131 field.clone(),
2132 2,
2133 Arc::new(Int32Array::from(vec![7, 8, 9, 10])),
2134 None,
2135 );
2136
2137 let result = interleave(&[&a, &b], &[(0, 2), (1, 0), (0, 0), (1, 1), (0, 1)]).unwrap();
2138 let result = result.as_fixed_size_list();
2139 assert_eq!(result.len(), 5);
2140 assert_eq!(result.value_length(), 2);
2141
2142 let values = result.values().as_primitive::<Int32Type>();
2143 assert_eq!(values.values(), &[5, 6, 7, 8, 1, 2, 9, 10, 3, 4]);
2145 }
2146
2147 #[test]
2148 fn test_interleave_fixed_size_list_with_nulls() {
2149 let field = Arc::new(Field::new("item", DataType::Int32, true));
2150 let a = FixedSizeListArray::new(
2152 field.clone(),
2153 2,
2154 Arc::new(Int32Array::from(vec![1, 2, 0, 0, 5, 6])),
2155 Some(NullBuffer::from(&[true, false, true])),
2156 );
2157 let b = FixedSizeListArray::new(
2159 field.clone(),
2160 2,
2161 Arc::new(Int32Array::from(vec![0, 0, 9, 10])),
2162 Some(NullBuffer::from(&[false, true])),
2163 );
2164
2165 let result = interleave(&[&a, &b], &[(0, 0), (0, 1), (1, 0), (1, 1), (0, 2)]).unwrap();
2166 let result = result.as_fixed_size_list();
2167 assert_eq!(result.len(), 5);
2168
2169 let validity: Vec<bool> = result.nulls().unwrap().iter().collect();
2170 assert_eq!(validity, &[true, false, false, true, true]);
2171 }
2172
2173 fn run_fsl_string_child_test(
2174 child_data_type: DataType,
2175 create_child: impl Fn(Vec<&str>) -> ArrayRef,
2176 extract_values: impl Fn(&ArrayRef) -> Vec<String>,
2177 ) {
2178 let field = Arc::new(Field::new("item", child_data_type, false));
2179
2180 let a = FixedSizeListArray::new(
2182 field.clone(),
2183 2,
2184 create_child(vec!["a", "b", "c", "d"]),
2185 None,
2186 );
2187 let b = FixedSizeListArray::new(
2189 field.clone(),
2190 2,
2191 create_child(vec!["x", "y", "z", "w"]),
2192 None,
2193 );
2194
2195 let result = interleave(&[&a, &b], &[(0, 1), (1, 0), (0, 0)]).unwrap();
2196 let result = result.as_fixed_size_list();
2197 assert_eq!(result.len(), 3);
2198 assert_eq!(result.value_length(), 2);
2199
2200 let values = extract_values(result.values());
2202 assert_eq!(values, vec!["c", "d", "x", "y", "a", "b"]);
2203 }
2204
2205 #[test]
2206 fn test_interleave_fixed_size_list_string_child() {
2207 run_fsl_string_child_test(
2209 DataType::Utf8,
2210 |v| Arc::new(StringArray::from(v)),
2211 |a| {
2212 a.as_string::<i32>()
2213 .iter()
2214 .map(|s| s.unwrap().to_owned())
2215 .collect()
2216 },
2217 );
2218 }
2219
2220 #[test]
2221 fn test_interleave_fixed_size_list_string_view_child() {
2222 run_fsl_string_child_test(
2224 DataType::Utf8View,
2225 |v| Arc::new(StringViewArray::from(v)),
2226 |a| {
2227 a.as_string_view()
2228 .iter()
2229 .map(|s| s.unwrap().to_owned())
2230 .collect()
2231 },
2232 );
2233 }
2234
2235 #[test]
2236 fn test_interleave_map() {
2237 use arrow_array::builder::MapBuilder;
2238 use arrow_array::builder::StringBuilder;
2239
2240 let mut a_builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
2242 a_builder.keys().append_value("k1");
2243 a_builder.values().append_value(1);
2244 a_builder.keys().append_value("k2");
2245 a_builder.values().append_value(2);
2246 a_builder.append(true).unwrap();
2247 a_builder.keys().append_value("k3");
2248 a_builder.values().append_value(3);
2249 a_builder.append(true).unwrap();
2250 let a = a_builder.finish();
2251
2252 let mut b_builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
2254 b_builder.keys().append_value("k4");
2255 b_builder.values().append_value(4);
2256 b_builder.append(true).unwrap();
2257 b_builder.keys().append_value("k5");
2258 b_builder.values().append_value(5);
2259 b_builder.keys().append_value("k6");
2260 b_builder.values().append_value(6);
2261 b_builder.keys().append_value("k7");
2262 b_builder.values().append_value(7);
2263 b_builder.append(true).unwrap();
2264 let b = b_builder.finish();
2265
2266 let result = interleave(&[&a, &b], &[(1, 0), (0, 0), (0, 1), (1, 1)]).unwrap();
2267 let result = result.as_map();
2268 assert_eq!(result.len(), 4);
2269
2270 let row0 = result.value(0);
2272 assert_eq!(row0.len(), 1);
2273 assert_eq!(row0.column(0).as_string::<i32>().value(0), "k4");
2274 assert_eq!(row0.column(1).as_primitive::<Int32Type>().value(0), 4);
2275
2276 let row1 = result.value(1);
2278 assert_eq!(row1.len(), 2);
2279
2280 let row2 = result.value(2);
2282 assert_eq!(row2.len(), 1);
2283 assert_eq!(row2.column(0).as_string::<i32>().value(0), "k3");
2284
2285 let row3 = result.value(3);
2287 assert_eq!(row3.len(), 3);
2288 }
2289
2290 #[test]
2291 fn test_interleave_map_with_nulls() {
2292 use arrow_array::builder::MapBuilder;
2293 use arrow_array::builder::StringBuilder;
2294
2295 let mut a_builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
2297 a_builder.keys().append_value("k1");
2298 a_builder.values().append_value(1);
2299 a_builder.append(true).unwrap();
2300 a_builder.append(false).unwrap();
2301 let a = a_builder.finish();
2302
2303 let mut b_builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
2305 b_builder.append(false).unwrap();
2306 b_builder.keys().append_value("k2");
2307 b_builder.values().append_value(2);
2308 b_builder.append(true).unwrap();
2309 let b = b_builder.finish();
2310
2311 let result = interleave(&[&a, &b], &[(0, 0), (1, 0), (0, 1), (1, 1)]).unwrap();
2312 let result = result.as_map();
2313 assert_eq!(result.len(), 4);
2314
2315 let validity: Vec<bool> = result.nulls().unwrap().iter().collect();
2316 assert_eq!(validity, &[true, false, false, true]);
2317
2318 assert_eq!(result.value(0).len(), 1);
2319 assert_eq!(result.value(3).len(), 1);
2320 }
2321}