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