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