1use std::fmt::Display;
21use std::mem::ManuallyDrop;
22use std::sync::Arc;
23
24use arrow_array::builder::UInt32Builder;
25use arrow_array::cast::AsArray;
26use arrow_array::types::*;
27use arrow_array::*;
28use arrow_buffer::{
29 ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer, NullBufferBuilder,
30 OffsetBuffer, RunEndBuffer, ScalarBuffer, bit_util,
31};
32use arrow_cmp::make_comparator;
33use arrow_data::{ArrayData, transform::MutableArrayData};
34use arrow_schema::{ArrowError, DataType, FieldRef, SortOptions, UnionFields, UnionMode};
35
36use num_traits::{CheckedAdd, Zero};
37
38pub fn take(
90 values: &dyn Array,
91 indices: &dyn Array,
92 options: Option<TakeOptions>,
93) -> Result<ArrayRef, ArrowError> {
94 let options = options.unwrap_or_default();
95 downcast_integer_array!(
96 indices => {
97 let indices = indices.to_indices();
98 if options.check_bounds {
99 check_bounds(values.len(), &indices)?;
100 }
101 take_impl::<_, true>(values, &indices)
102 },
103 d => Err(ArrowError::InvalidArgumentError(format!("Take only supported for integers, got {d:?}")))
104 )
105}
106
107pub fn take_arrays(
156 arrays: &[ArrayRef],
157 indices: &dyn Array,
158 options: Option<TakeOptions>,
159) -> Result<Vec<ArrayRef>, ArrowError> {
160 arrays
161 .iter()
162 .map(|array| take(array.as_ref(), indices, options.clone()))
163 .collect()
164}
165
166pub unsafe fn take_arrays_unchecked(
178 arrays: &[ArrayRef],
179 indices: &dyn Array,
180) -> Result<Vec<ArrayRef>, ArrowError> {
181 downcast_integer_array!(
182 indices => {
183 let indices = indices.to_indices();
184 arrays
185 .iter()
186 .map(|array| take_impl::<_, false>(array.as_ref(), &indices))
187 .collect()
188 },
189 d => Err(ArrowError::InvalidArgumentError(format!("Take only supported for integers, got {d:?}")))
190 )
191}
192
193fn check_bounds<T: ArrowPrimitiveType>(
195 len: usize,
196 indices: &PrimitiveArray<T>,
197) -> Result<(), ArrowError>
198where
199 T::Native: Display,
200{
201 let Some(len) = T::Native::from_usize(len) else {
202 return if T::DATA_TYPE.is_integer() {
203 Ok(())
205 } else {
206 Err(ArrowError::ComputeError("Cast to usize failed".to_string()))
207 };
208 };
209
210 if indices.null_count() > 0 {
211 indices.iter().flatten().try_for_each(|index| {
212 if index >= len {
213 return Err(ArrowError::ComputeError(format!(
214 "Array index out of bounds, cannot get item at index {index} from {len} entries"
215 )));
216 }
217 Ok(())
218 })
219 } else {
220 let in_bounds = indices.values().iter().fold(true, |in_bounds, &i| {
221 in_bounds & (i >= T::Native::ZERO) & (i < len)
222 });
223
224 if !in_bounds {
225 for &index in indices.values() {
226 if index < T::Native::ZERO || index >= len {
227 return Err(ArrowError::ComputeError(format!(
228 "Array index out of bounds, cannot get item at index {index} from {len} entries"
229 )));
230 }
231 }
232 }
233
234 Ok(())
235 }
236}
237
238#[inline(never)]
239fn take_impl<IndexType: ArrowPrimitiveType, const VALIDATE_INDICES: bool>(
240 values: &dyn Array,
241 indices: &PrimitiveArray<IndexType>,
242) -> Result<ArrayRef, ArrowError> {
243 if indices.is_empty() {
244 if let DataType::Union(fields, _) = values.data_type()
245 && fields.is_empty()
246 {
247 return Ok(values.slice(0, 0));
250 }
251 return Ok(new_empty_array(values.data_type()));
252 }
253 downcast_primitive_array! {
254 values => Ok(Arc::new(take_primitive::<_, _, VALIDATE_INDICES>(values, indices)?)),
255 DataType::Boolean => {
256 let values = values.as_any().downcast_ref::<BooleanArray>().unwrap();
257 Ok(Arc::new(take_boolean::<_, VALIDATE_INDICES>(values, indices)))
258 }
259 DataType::Utf8 => {
260 Ok(Arc::new(take_bytes::<_, _, VALIDATE_INDICES>(values.as_string::<i32>(), indices)?))
261 }
262 DataType::LargeUtf8 => {
263 Ok(Arc::new(take_bytes::<_, _, VALIDATE_INDICES>(values.as_string::<i64>(), indices)?))
264 }
265 DataType::Utf8View => {
266 Ok(Arc::new(take_byte_view::<_, _, VALIDATE_INDICES>(values.as_string_view(), indices)?))
267 }
268 DataType::List(_) => {
269 Ok(Arc::new(take_list::<_, Int32Type, VALIDATE_INDICES>(values.as_list(), indices)?))
270 }
271 DataType::LargeList(_) => {
272 Ok(Arc::new(take_list::<_, Int64Type, VALIDATE_INDICES>(values.as_list(), indices)?))
273 }
274 DataType::ListView(_) => {
275 Ok(Arc::new(take_list_view::<_, Int32Type, VALIDATE_INDICES>(values.as_list_view(), indices)?))
276 }
277 DataType::LargeListView(_) => {
278 Ok(Arc::new(take_list_view::<_, Int64Type, VALIDATE_INDICES>(values.as_list_view(), indices)?))
279 }
280 DataType::FixedSizeList(_, length) => {
281 let values = values
282 .as_any()
283 .downcast_ref::<FixedSizeListArray>()
284 .unwrap();
285 Ok(Arc::new(take_fixed_size_list::<_, VALIDATE_INDICES>(
286 values,
287 indices,
288 *length as u32,
289 )?))
290 }
291 DataType::Map(field, ordered) => {
292 let list_arr = ListArray::from(values.as_map().clone());
293 let list_data = take_list::<_, Int32Type, VALIDATE_INDICES>(&list_arr, indices)?;
294 let (_, offsets, entries, nulls) = list_data.into_parts();
295 let entries = entries.as_struct().clone();
296 Ok(Arc::new(MapArray::try_new(
297 field.clone(),
298 offsets,
299 entries,
300 nulls,
301 *ordered,
302 )?))
303 }
304 DataType::Struct(fields) => {
305 let array: &StructArray = values.as_struct();
306 let arrays = array
307 .columns()
308 .iter()
309 .map(|a| take_impl::<_, VALIDATE_INDICES>(a.as_ref(), indices))
310 .collect::<Result<Vec<ArrayRef>, _>>()?;
311 let fields: Vec<(FieldRef, ArrayRef)> =
312 fields.iter().cloned().zip(arrays).collect();
313
314 let is_valid: Buffer = indices
316 .iter()
317 .map(|index| {
318 if let Some(index) = index {
319 array.is_valid(index.to_usize().unwrap())
320 } else {
321 false
322 }
323 })
324 .collect();
325
326 if fields.is_empty() {
327 let nulls = NullBuffer::new(BooleanBuffer::new(is_valid, 0, indices.len()));
328 Ok(Arc::new(StructArray::new_empty_fields(indices.len(), Some(nulls))))
329 } else {
330 Ok(Arc::new(StructArray::from((fields, is_valid))) as ArrayRef)
331 }
332 }
333 DataType::Dictionary(_, _) => downcast_dictionary_array! {
334 values => Ok(Arc::new(take_dict::<_, _, VALIDATE_INDICES>(values, indices)?)),
335 t => unimplemented!("Take not supported for dictionary type {:?}", t)
336 }
337 DataType::RunEndEncoded(_, _) => downcast_run_array! {
338 values => Ok(Arc::new(take_run(values, indices)?)),
339 t => unimplemented!("Take not supported for run type {:?}", t)
340 }
341 DataType::Binary => {
342 Ok(Arc::new(take_bytes::<_, _, VALIDATE_INDICES>(values.as_binary::<i32>(), indices)?))
343 }
344 DataType::LargeBinary => {
345 Ok(Arc::new(take_bytes::<_, _, VALIDATE_INDICES>(values.as_binary::<i64>(), indices)?))
346 }
347 DataType::BinaryView => {
348 Ok(Arc::new(take_byte_view::<_, _, VALIDATE_INDICES>(values.as_binary_view(), indices)?))
349 }
350 DataType::FixedSizeBinary(size) => {
351 let values = values
352 .as_any()
353 .downcast_ref::<FixedSizeBinaryArray>()
354 .unwrap();
355 Ok(Arc::new(take_fixed_size_binary::<_, VALIDATE_INDICES>(values, indices, *size)?))
356 }
357 DataType::Null => {
358 if values.len() >= indices.len() {
360 Ok(values.slice(0, indices.len()))
363 } else {
364 Ok(new_null_array(&DataType::Null, indices.len()))
366 }
367 }
368 DataType::Union(fields, UnionMode::Sparse) => {
369 let values = values.as_any().downcast_ref::<UnionArray>().unwrap();
370 let (type_ids, null_type_id) =
371 take_union_type_ids(fields, values.type_ids(), indices)?;
372 let children = fields
373 .iter()
374 .map(|(type_id, _)| {
375 take_sparse_union_child::<_, VALIDATE_INDICES>(
376 values.child(type_id),
377 null_type_id == Some(type_id),
378 indices,
379 )
380 })
381 .collect::<Result<Vec<_>, _>>()?;
382 let array = UnionArray::try_new(fields.clone(), type_ids, None, children)?;
383 Ok(Arc::new(array))
384 }
385 DataType::Union(fields, UnionMode::Dense) => {
386 let values = values.as_any().downcast_ref::<UnionArray>().unwrap();
387
388 let (type_ids, _) = take_union_type_ids(fields, values.type_ids(), indices)?;
389 let type_ids = PrimitiveArray::<Int8Type>::try_new(type_ids, None)?;
390 let offsets = <PrimitiveArray<Int32Type>>::try_new(
393 take_native::<_, _, VALIDATE_INDICES>(values.offsets().unwrap(), indices),
394 indices.nulls().cloned(),
395 )?;
396
397 let children = fields.iter()
398 .map(|(field_type_id, _)| {
399 let mask = BooleanArray::from_unary(&type_ids, |value_type_id| value_type_id == field_type_id);
400
401 let indices = crate::filter::filter(&offsets, &mask)?;
402
403 let values = values.child(field_type_id);
404
405 take_impl::<_, VALIDATE_INDICES>(values, indices.as_primitive::<Int32Type>())
406 })
407 .collect::<Result<_, _>>()?;
408
409 let mut child_offsets = [0; 128];
410
411 let offsets = type_ids.values()
412 .iter()
413 .map(|&i| {
414 let offset = child_offsets[i as usize];
415
416 child_offsets[i as usize] += 1;
417
418 offset
419 })
420 .collect();
421
422 let (_, type_ids, _) = type_ids.into_parts();
423
424 let array = UnionArray::try_new(fields.clone(), type_ids, Some(offsets), children)?;
425
426 Ok(Arc::new(array))
427 }
428 t => unimplemented!("Take not supported for data type {:?}", t)
429 }
430}
431
432fn take_union_type_ids<IndexType: ArrowPrimitiveType>(
443 fields: &UnionFields,
444 type_ids: &ScalarBuffer<i8>,
445 indices: &PrimitiveArray<IndexType>,
446) -> Result<(ScalarBuffer<i8>, Option<i8>), ArrowError> {
447 if indices.null_count() == 0 {
448 return Ok((take_native::<_, _, true>(type_ids, indices), None));
449 }
450
451 let null_type_id = union_null_type_id(fields)?;
452 let taken_type_ids = take_native::<_, _, true>(type_ids, indices);
453 let type_ids = indices
454 .iter()
455 .zip(&taken_type_ids)
456 .map(|(index, &type_id)| {
457 if index.is_some() {
458 type_id
459 } else {
460 null_type_id
461 }
462 })
463 .collect::<ScalarBuffer<_>>();
464 Ok((type_ids, Some(null_type_id)))
465}
466
467fn union_null_type_id(fields: &UnionFields) -> Result<i8, ArrowError> {
469 fields
470 .iter()
471 .find_map(|(type_id, field)| field_can_represent_take_null(field).then_some(type_id))
472 .ok_or_else(|| {
473 ArrowError::ComputeError(
474 "Cannot take null indices from a union with no field that can represent nulls"
475 .into(),
476 )
477 })
478}
479
480fn field_can_represent_take_null(field: &FieldRef) -> bool {
486 if !field.is_nullable() {
487 return false;
488 }
489 match field.data_type() {
490 DataType::RunEndEncoded(_, values) => field_can_represent_take_null(values),
491 DataType::Union(fields, _) => fields
492 .iter()
493 .any(|(_, child)| field_can_represent_take_null(child)),
494 _ => true,
495 }
496}
497
498fn take_sparse_union_child<IndexType: ArrowPrimitiveType, const VALIDATE_INDICES: bool>(
505 values: &dyn Array,
506 represent_nulls: bool,
507 indices: &PrimitiveArray<IndexType>,
508) -> Result<ArrayRef, ArrowError> {
509 if represent_nulls || indices.null_count() == 0 {
510 return take_impl::<_, VALIDATE_INDICES>(values, indices);
511 }
512
513 if values.is_empty() {
514 return Ok(make_array(
518 new_null_array(values.data_type(), indices.len())
519 .to_data()
520 .into_builder()
521 .nulls(None)
522 .build()?,
523 ));
524 }
525
526 take_impl::<_, VALIDATE_INDICES>(values, &indices_without_nulls(indices))
527}
528
529fn indices_without_nulls<IndexType: ArrowPrimitiveType>(
531 indices: &PrimitiveArray<IndexType>,
532) -> PrimitiveArray<IndexType> {
533 let dummy = IndexType::Native::ZERO;
534 let normalized = indices.iter().map(|idx| idx.unwrap_or(dummy));
535 PrimitiveArray::from_iter_values(normalized)
536}
537
538#[derive(Clone, Debug, Default)]
540pub struct TakeOptions {
541 pub check_bounds: bool,
545}
546
547fn take_primitive<T, I, const VALIDATE_INDICES: bool>(
557 values: &PrimitiveArray<T>,
558 indices: &PrimitiveArray<I>,
559) -> Result<PrimitiveArray<T>, ArrowError>
560where
561 T: ArrowPrimitiveType,
562 I: ArrowPrimitiveType,
563{
564 let values_buf = take_native::<_, _, VALIDATE_INDICES>(values.values(), indices);
565 let nulls = take_nulls::<_, VALIDATE_INDICES>(values.nulls(), indices);
566 Ok(PrimitiveArray::try_new(values_buf, nulls)?.with_data_type(values.data_type().clone()))
567}
568
569#[inline(never)]
570fn take_nulls<I: ArrowPrimitiveType, const VALIDATE_INDICES: bool>(
571 values: Option<&NullBuffer>,
572 indices: &PrimitiveArray<I>,
573) -> Option<NullBuffer> {
574 match values.filter(|n| n.null_count() > 0) {
575 Some(n) => NullBuffer::from_unsliced_buffer(
576 take_bits::<_, VALIDATE_INDICES>(n.inner(), indices).into_inner(),
577 indices.len(),
578 ),
579 None => indices.nulls().cloned(),
580 }
581}
582
583#[inline(never)]
584fn take_native<T: ArrowNativeType, I: ArrowPrimitiveType, const VALIDATE_INDICES: bool>(
585 values: &[T],
586 indices: &PrimitiveArray<I>,
587) -> ScalarBuffer<T> {
588 match indices.nulls().filter(|n| n.null_count() > 0) {
589 Some(n) => indices
590 .values()
591 .iter()
592 .enumerate()
593 .map(|(idx, index)| match values.get(index.as_usize()) {
594 Some(v) => *v,
595 None => match unsafe { n.inner().value_unchecked(idx) } {
597 false => T::default(),
598 true => panic!("Out-of-bounds index {index:?}"),
599 },
600 })
601 .collect(),
602 None => {
603 if VALIDATE_INDICES {
604 indices
605 .values()
606 .iter()
607 .map(|index| values[index.as_usize()])
608 .collect()
609 } else {
610 indices
611 .values()
612 .iter()
613 .map(|index| {
614 unsafe { *values.get_unchecked(index.as_usize()) }
616 })
617 .collect()
618 }
619 }
620 }
621}
622
623#[inline(always)]
635unsafe fn copy_bit_if_set(src: *const u8, src_bit_idx: usize, dst: *mut u8, dst_bit_idx: usize) {
636 unsafe {
637 if bit_util::get_bit_raw(src, src_bit_idx) {
638 bit_util::set_bit_raw(dst, dst_bit_idx);
639 }
640 }
641}
642
643#[inline(always)]
654unsafe fn pack_bit(src: *const u8, bit_idx: usize, out_pos: usize) -> u8 {
655 let byte = unsafe { *src.add(bit_idx >> 3) }; ((byte >> (bit_idx & 7)) & 1) << out_pos }
658
659#[inline(never)]
660fn take_bits<I: ArrowPrimitiveType, const VALIDATE_INDICES: bool>(
661 values: &BooleanBuffer,
662 indices: &PrimitiveArray<I>,
663) -> BooleanBuffer {
664 let len = indices.len();
665 let src_offset = values.offset();
666 let src_ptr = values.values().as_ptr();
667 let out_bytes = len.div_ceil(8);
668
669 match indices.nulls().filter(|nulls| nulls.null_count() > 0) {
670 Some(index_nulls) => {
671 let mut output = vec![0u8; out_bytes];
672 let out_ptr = output.as_mut_ptr();
673 index_nulls.valid_indices().for_each(|valid_idx| {
674 let index_val = unsafe { indices.value_unchecked(valid_idx) }.as_usize();
676 if VALIDATE_INDICES {
677 if values.value(index_val) {
678 unsafe { bit_util::set_bit_raw(out_ptr, valid_idx) };
680 }
681 } else {
682 unsafe { copy_bit_if_set(src_ptr, index_val + src_offset, out_ptr, valid_idx) };
684 }
685 });
686 BooleanBuffer::new(Buffer::from(output), 0, len)
687 }
688 None => {
689 let mut output = vec![0u8; out_bytes];
692 let out_slice = output.as_mut_slice();
693 let full_bytes = len / 8;
694
695 for (byte_idx, out_byte) in out_slice.iter_mut().enumerate().take(full_bytes) {
696 let base = byte_idx * 8;
697 let mut byte = 0u8;
698 for bit in 0..8usize {
699 let index_val = unsafe { indices.value_unchecked(base + bit) }.as_usize();
702 if VALIDATE_INDICES {
703 byte |= (values.value(index_val) as u8) << bit;
704 } else {
705 byte |= unsafe { pack_bit(src_ptr, index_val + src_offset, bit) };
707 }
708 }
709 *out_byte = byte;
710 }
711 if full_bytes < out_bytes {
713 let base = full_bytes * 8;
714 let mut byte = 0u8;
715 for bit in 0..(len - base) {
716 let index_val = unsafe { indices.value_unchecked(base + bit) }.as_usize();
719 if VALIDATE_INDICES {
720 byte |= (values.value(index_val) as u8) << bit;
721 } else {
722 byte |= unsafe { pack_bit(src_ptr, index_val + src_offset, bit) };
724 }
725 }
726 out_slice[full_bytes] = byte;
727 }
728 BooleanBuffer::new(Buffer::from(output), 0, len)
729 }
730 }
731}
732
733#[inline(never)]
736fn take_bits_with_validity<I: ArrowPrimitiveType, const VALIDATE_INDICES: bool>(
737 values: &BooleanBuffer,
738 validity: &BooleanBuffer,
739 indices: &PrimitiveArray<I>,
740) -> (BooleanBuffer, Option<NullBuffer>) {
741 let len = indices.len();
742 let value_bit_offset = values.offset();
743 let validity_bit_offset = validity.offset();
744 let value_data_ptr = values.values().as_ptr();
745 let validity_data_ptr = validity.values().as_ptr();
746 let out_bytes = len.div_ceil(8);
747
748 let mut value_out = vec![0u8; out_bytes];
749 let mut validity_out = vec![0u8; out_bytes];
750
751 match indices.nulls().filter(|nulls| nulls.null_count() > 0) {
752 Some(index_nulls) => {
753 let value_out_ptr = value_out.as_mut_ptr();
754 let validity_out_ptr = validity_out.as_mut_ptr();
755 for out_pos in index_nulls.valid_indices() {
756 let src_idx = unsafe { indices.value_unchecked(out_pos) }.as_usize();
758 if VALIDATE_INDICES {
759 if values.value(src_idx) {
760 unsafe { bit_util::set_bit_raw(value_out_ptr, out_pos) };
762 }
763 if validity.value(src_idx) {
764 unsafe { bit_util::set_bit_raw(validity_out_ptr, out_pos) };
766 }
767 } else {
768 unsafe {
770 copy_bit_if_set(
771 value_data_ptr,
772 src_idx + value_bit_offset,
773 value_out_ptr,
774 out_pos,
775 );
776 copy_bit_if_set(
777 validity_data_ptr,
778 src_idx + validity_bit_offset,
779 validity_out_ptr,
780 out_pos,
781 );
782 }
783 }
784 }
785 }
786 None => {
787 let value_out_slice = value_out.as_mut_slice();
788 let validity_out_slice = validity_out.as_mut_slice();
789 let full_bytes = len / 8;
790
791 for (byte_idx, (value_out_byte, validity_out_byte)) in value_out_slice
792 .iter_mut()
793 .zip(validity_out_slice.iter_mut())
794 .enumerate()
795 .take(full_bytes)
796 {
797 let bit_base = byte_idx * 8;
798 let mut packed_values = 0u8;
799 let mut packed_validity = 0u8;
800 for bit_pos in 0..8usize {
801 let src_idx = unsafe { indices.value_unchecked(bit_base + bit_pos) }.as_usize();
803 if VALIDATE_INDICES {
804 packed_values |= (values.value(src_idx) as u8) << bit_pos;
805 packed_validity |= (validity.value(src_idx) as u8) << bit_pos;
806 } else {
807 packed_values |= unsafe {
809 pack_bit(value_data_ptr, src_idx + value_bit_offset, bit_pos)
810 };
811 packed_validity |= unsafe {
812 pack_bit(validity_data_ptr, src_idx + validity_bit_offset, bit_pos)
813 };
814 }
815 }
816 *value_out_byte = packed_values;
817 *validity_out_byte = packed_validity;
818 }
819 if full_bytes < out_bytes {
821 let bit_base = full_bytes * 8;
822 let mut packed_values = 0u8;
823 let mut packed_validity = 0u8;
824 for bit_pos in 0..(len - bit_base) {
825 let src_idx = unsafe { indices.value_unchecked(bit_base + bit_pos) }.as_usize();
827 if VALIDATE_INDICES {
828 packed_values |= (values.value(src_idx) as u8) << bit_pos;
829 packed_validity |= (validity.value(src_idx) as u8) << bit_pos;
830 } else {
831 packed_values |= unsafe {
833 pack_bit(value_data_ptr, src_idx + value_bit_offset, bit_pos)
834 };
835 packed_validity |= unsafe {
836 pack_bit(validity_data_ptr, src_idx + validity_bit_offset, bit_pos)
837 };
838 }
839 }
840 value_out_slice[full_bytes] = packed_values;
841 validity_out_slice[full_bytes] = packed_validity;
842 }
843 }
844 }
845
846 let value_buf_out = BooleanBuffer::new(Buffer::from(value_out), 0, len);
847 let validity_buf_out = NullBuffer::from_unsliced_buffer(validity_out, len);
848 (value_buf_out, validity_buf_out)
849}
850
851fn take_boolean<IndexType: ArrowPrimitiveType, const VALIDATE_INDICES: bool>(
853 array: &BooleanArray,
854 indices: &PrimitiveArray<IndexType>,
855) -> BooleanArray {
856 let bits = array.values();
857 match array.nulls().filter(|n| n.null_count() > 0) {
858 Some(array_nulls) => {
859 let (val_buf, null_buf) =
860 take_bits_with_validity::<_, VALIDATE_INDICES>(bits, array_nulls.inner(), indices);
861 BooleanArray::new(val_buf, null_buf)
862 }
863 None => {
864 let val_buf = take_bits::<_, VALIDATE_INDICES>(bits, indices);
865 let null_buf = take_nulls::<_, VALIDATE_INDICES>(None, indices);
866 BooleanArray::new(val_buf, null_buf)
867 }
868 }
869}
870
871fn take_bytes<T: ByteArrayType, IndexType: ArrowPrimitiveType, const VALIDATE_INDICES: bool>(
873 array: &GenericByteArray<T>,
874 indices: &PrimitiveArray<IndexType>,
875) -> Result<GenericByteArray<T>, ArrowError> {
876 let mut values: Vec<u8> = Vec::new();
877 let mut offsets = Vec::with_capacity(indices.len() + 1);
878 offsets.push(T::Offset::default());
879
880 let input_offsets = array.value_offsets();
881 let mut capacity = 0;
882 let nulls = take_nulls::<_, VALIDATE_INDICES>(array.nulls(), indices);
883
884 match nulls.as_ref().filter(|n| n.null_count() > 0) {
886 None => {
888 for index in indices.values() {
889 let index = index.as_usize();
890 let start = input_offsets[index].as_usize();
891 let end = input_offsets[index + 1].as_usize();
892 capacity += end - start;
893 offsets.push(
894 T::Offset::from_usize(capacity)
895 .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?,
896 );
897 }
898
899 values.reserve(capacity);
900
901 let dst = values.spare_capacity_mut();
902 debug_assert!(dst.len() >= capacity);
903 let mut offset = 0;
904
905 for index in indices.values() {
906 unsafe {
909 let data: &[u8] = array.value_unchecked(index.as_usize()).as_ref();
910 std::ptr::copy_nonoverlapping(
911 data.as_ptr(),
912 dst.get_unchecked_mut(offset..).as_mut_ptr().cast::<u8>(),
913 data.len(),
914 );
915 offset += data.len();
916 }
917 }
918
919 unsafe {
921 values.set_len(capacity);
922 }
923 }
924 Some(output_nulls) => {
926 let mut source_ranges = Vec::with_capacity(indices.len() - output_nulls.null_count());
927 let mut last_filled = 0;
928
929 offsets.resize(indices.len() + 1, T::Offset::default());
931
932 for i in output_nulls.valid_indices() {
934 let current_offset = T::Offset::from_usize(capacity)
935 .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?;
936 if last_filled < i {
938 offsets[last_filled + 1..=i].fill(current_offset);
939 }
940
941 let index = unsafe { indices.value_unchecked(i) }.as_usize();
943 let start = input_offsets[index].as_usize();
944 let end = input_offsets[index + 1].as_usize();
945 capacity += end - start;
946 offsets[i + 1] = T::Offset::from_usize(capacity)
947 .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?;
948
949 source_ranges.push((start, end));
950 last_filled = i + 1;
951 }
952
953 let final_offset = T::Offset::from_usize(capacity)
955 .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?;
956 offsets[last_filled + 1..].fill(final_offset);
957 values.reserve(capacity);
959 debug_assert_eq!(
960 source_ranges.iter().map(|(s, e)| e - s).sum::<usize>(),
961 capacity,
962 "capacity must equal total bytes across all ranges"
963 );
964
965 let src = array.value_data();
966 let src = src.as_ptr();
967 let dst = values.spare_capacity_mut();
968 debug_assert!(dst.len() >= capacity);
969
970 let mut offset = 0;
971
972 for (start, end) in source_ranges {
973 let value_len = end - start;
974 unsafe {
978 std::ptr::copy_nonoverlapping(
979 src.add(start),
980 dst.get_unchecked_mut(offset..).as_mut_ptr().cast::<u8>(),
981 value_len,
982 );
983 offset += value_len;
984 }
985 }
986 unsafe { values.set_len(capacity) };
989 }
990 }
991
992 let array = unsafe {
995 let offsets = OffsetBuffer::new_unchecked(offsets.into());
996 GenericByteArray::<T>::new_unchecked(offsets, values.into(), nulls)
997 };
998
999 Ok(array)
1000}
1001
1002fn take_byte_view<T: ByteViewType, IndexType: ArrowPrimitiveType, const VALIDATE_INDICES: bool>(
1004 array: &GenericByteViewArray<T>,
1005 indices: &PrimitiveArray<IndexType>,
1006) -> Result<GenericByteViewArray<T>, ArrowError> {
1007 let new_views = take_native::<_, _, VALIDATE_INDICES>(array.views(), indices);
1008 let new_nulls = take_nulls::<_, VALIDATE_INDICES>(array.nulls(), indices);
1009 let buffers = Arc::clone(array.data_buffers());
1010 Ok(unsafe { GenericByteViewArray::new_unchecked(new_views, buffers, new_nulls) })
1012}
1013
1014fn take_list<IndexType, OffsetType, const VALIDATE_INDICES: bool>(
1019 values: &GenericListArray<OffsetType::Native>,
1020 indices: &PrimitiveArray<IndexType>,
1021) -> Result<GenericListArray<OffsetType::Native>, ArrowError>
1022where
1023 IndexType: ArrowPrimitiveType,
1024 OffsetType: ArrowPrimitiveType,
1025 OffsetType::Native: OffsetSizeTrait,
1026 PrimitiveArray<OffsetType>: From<Vec<OffsetType::Native>>,
1027{
1028 let src_offsets = values.value_offsets();
1029 let child_data = values.values().to_data();
1030 let visible_child_len = (values.offsets().last() - values.offsets().first()).as_usize();
1032 let estimated_child_len = visible_child_len
1034 .checked_div(values.len())
1035 .map_or(0, |avg| avg.saturating_mul(indices.len()));
1036 let nulls = take_nulls::<_, VALIDATE_INDICES>(values.nulls(), indices);
1037
1038 let mut dst_offsets = Vec::with_capacity(indices.len() + 1);
1039 dst_offsets.push(OffsetType::Native::zero());
1040
1041 let field = values.value_field().clone();
1042
1043 if child_data.null_count() == 0
1044 && let Some(bytes_per_value) = child_data.data_type().primitive_width()
1045 {
1046 let values_buf = &child_data.buffers()[0];
1047 let child_buf_offset = child_data.offset() * bytes_per_value;
1048
1049 let mut dst_buf = MutableBuffer::new(estimated_child_len.saturating_mul(bytes_per_value));
1050
1051 let mut child_len = OffsetType::Native::zero();
1052
1053 match nulls.as_ref().filter(|n| n.null_count() > 0) {
1054 None => {
1055 for &idx in indices.values() {
1056 let row = idx.as_usize();
1057 let start = child_buf_offset + src_offsets[row].as_usize() * bytes_per_value;
1058 let end = child_buf_offset + src_offsets[row + 1].as_usize() * bytes_per_value;
1059 dst_buf.extend_from_slice(&values_buf[start..end]);
1060 child_len = child_len
1061 .checked_add(&(src_offsets[row + 1] - src_offsets[row]))
1062 .ok_or_else(|| ArrowError::OffsetOverflowError(child_len.as_usize()))?;
1063 dst_offsets.push(child_len);
1064 }
1065 }
1066 Some(valid) => {
1067 let mut prev = 0;
1068 for vidx in valid.valid_indices() {
1069 if prev < vidx {
1071 dst_offsets.extend(std::iter::repeat_n(child_len, vidx - prev));
1072 }
1073 let row = unsafe { indices.value_unchecked(vidx) }.as_usize();
1075 let (start, end) = if VALIDATE_INDICES {
1076 (
1077 child_buf_offset + src_offsets[row].as_usize() * bytes_per_value,
1078 child_buf_offset + src_offsets[row + 1].as_usize() * bytes_per_value,
1079 )
1080 } else {
1081 unsafe {
1083 (
1084 child_buf_offset
1085 + src_offsets.get_unchecked(row).as_usize() * bytes_per_value,
1086 child_buf_offset
1087 + src_offsets.get_unchecked(row + 1).as_usize()
1088 * bytes_per_value,
1089 )
1090 }
1091 };
1092 dst_buf.extend_from_slice(&values_buf[start..end]);
1093 child_len = child_len
1094 .checked_add(&(src_offsets[row + 1] - src_offsets[row]))
1095 .ok_or_else(|| ArrowError::OffsetOverflowError(child_len.as_usize()))?;
1096 dst_offsets.push(child_len);
1097 prev = vidx + 1;
1098 }
1099 dst_offsets.extend(std::iter::repeat_n(child_len, indices.len() - prev));
1100 }
1101 }
1102
1103 debug_assert_eq!(
1104 dst_offsets.len(),
1105 indices.len() + 1,
1106 "New offsets was filled under/over the expected capacity"
1107 );
1108
1109 let child = make_array(unsafe {
1112 ArrayData::builder(child_data.data_type().clone())
1113 .len(child_len.as_usize())
1114 .add_buffer(dst_buf.into())
1115 .build_unchecked()
1116 });
1117 let offsets = unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(dst_offsets)) };
1119 return GenericListArray::<OffsetType::Native>::try_new(field, offsets, child, nulls);
1120 }
1121
1122 let mut mutable = MutableArrayData::new(
1123 vec![&child_data],
1124 child_data.null_count() > 0,
1125 estimated_child_len,
1126 );
1127
1128 match nulls.as_ref().filter(|n| n.null_count() > 0) {
1129 None => {
1130 for idx in indices.values() {
1131 let row = idx.as_usize();
1132 mutable.try_extend(
1133 0,
1134 src_offsets[row].as_usize(),
1135 src_offsets[row + 1].as_usize(),
1136 )?;
1137 dst_offsets.push(
1138 OffsetType::Native::from_usize(mutable.len())
1139 .ok_or_else(|| ArrowError::OffsetOverflowError(mutable.len()))?,
1140 );
1141 }
1142 }
1143 Some(valid) => {
1144 let mut last = 0;
1145 for i in valid.valid_indices() {
1146 let current = OffsetType::Native::from_usize(mutable.len())
1147 .ok_or_else(|| ArrowError::OffsetOverflowError(mutable.len()))?;
1148 if last < i {
1149 dst_offsets.extend(std::iter::repeat_n(current, i - last));
1150 }
1151 let row = unsafe { indices.value_unchecked(i) }.as_usize();
1153 let (start, end) = if VALIDATE_INDICES {
1154 (src_offsets[row].as_usize(), src_offsets[row + 1].as_usize())
1155 } else {
1156 unsafe {
1158 (
1159 src_offsets.get_unchecked(row).as_usize(),
1160 src_offsets.get_unchecked(row + 1).as_usize(),
1161 )
1162 }
1163 };
1164 mutable.try_extend(0, start, end)?;
1165 dst_offsets.push(
1166 OffsetType::Native::from_usize(mutable.len())
1167 .ok_or_else(|| ArrowError::OffsetOverflowError(mutable.len()))?,
1168 );
1169 last = i + 1;
1170 }
1171 let final_offset = OffsetType::Native::from_usize(mutable.len())
1173 .ok_or_else(|| ArrowError::OffsetOverflowError(mutable.len()))?;
1174 dst_offsets.extend(std::iter::repeat_n(final_offset, indices.len() - last));
1175 }
1176 }
1177
1178 debug_assert_eq!(dst_offsets.len(), indices.len() + 1);
1179
1180 let offsets = unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(dst_offsets)) };
1182 let child = make_array(mutable.freeze());
1183 GenericListArray::<OffsetType::Native>::try_new(field, offsets, child, nulls)
1184}
1185
1186fn take_list_view<IndexType, OffsetType, const VALIDATE_INDICES: bool>(
1187 values: &GenericListViewArray<OffsetType::Native>,
1188 indices: &PrimitiveArray<IndexType>,
1189) -> Result<GenericListViewArray<OffsetType::Native>, ArrowError>
1190where
1191 IndexType: ArrowPrimitiveType,
1192 OffsetType: ArrowPrimitiveType,
1193 OffsetType::Native: OffsetSizeTrait,
1194{
1195 let taken_offsets = take_native::<_, _, VALIDATE_INDICES>(values.offsets(), indices);
1196 let taken_sizes = take_native::<_, _, VALIDATE_INDICES>(values.sizes(), indices);
1197 let nulls = take_nulls::<_, VALIDATE_INDICES>(values.nulls(), indices);
1198
1199 let field = match values.data_type() {
1200 DataType::ListView(field) | DataType::LargeListView(field) => field.clone(),
1201 d => unreachable!("take_list_view called with non-list-view data type {d}"),
1202 };
1203
1204 Ok(unsafe {
1207 GenericListViewArray::<OffsetType::Native>::new_unchecked(
1208 field,
1209 taken_offsets,
1210 taken_sizes,
1211 Arc::clone(values.values()),
1212 nulls,
1213 )
1214 })
1215}
1216
1217fn take_fixed_size_list<IndexType: ArrowPrimitiveType, const VALIDATE_INDICES: bool>(
1223 values: &FixedSizeListArray,
1224 indices: &PrimitiveArray<IndexType>,
1225 length: <UInt32Type as ArrowPrimitiveType>::Native,
1226) -> Result<FixedSizeListArray, ArrowError> {
1227 let field = values.value_field();
1228 let child = values.values();
1229 let nulls = take_nulls::<_, VALIDATE_INDICES>(values.nulls(), indices);
1230
1231 let taken_child = if child.null_count() == 0
1233 && let Some(element_size) = child.data_type().primitive_width()
1234 {
1235 take_fixed_size_list_primitive(
1236 child,
1237 indices,
1238 length as usize,
1239 element_size,
1240 nulls.as_ref(),
1241 )
1242 } else {
1243 let list_indices = take_value_indices_from_fixed_size_list(values, indices, length)?;
1244 take_impl::<UInt32Type, VALIDATE_INDICES>(child.as_ref(), &list_indices)?
1245 };
1246
1247 FixedSizeListArray::try_new_with_length(
1248 field.clone(),
1249 length as i32,
1250 taken_child,
1251 nulls,
1252 indices.len(),
1253 )
1254}
1255
1256#[inline(never)]
1257fn take_fixed_size_list_primitive<IndexType: ArrowPrimitiveType>(
1258 child: &ArrayRef,
1259 indices: &PrimitiveArray<IndexType>,
1260 list_size: usize,
1261 element_size: usize,
1262 taken_row_nulls: Option<&NullBuffer>,
1263) -> ArrayRef {
1264 let row_bytes = list_size * element_size;
1265 let child_data = child.to_data();
1266 let src = child_data.buffers()[0].as_slice();
1267 let child_byte_offset = child_data.offset() * element_size;
1268
1269 debug_assert!(
1270 indices.len().checked_mul(row_bytes).is_some(),
1271 "take_fixed_size_list_primitive: output buffer size overflows usize"
1272 );
1273 let out_len = indices.len() * list_size;
1274
1275 let mut out = MutableBuffer::from_len_zeroed(indices.len() * row_bytes);
1276 let out_slice = out.as_slice_mut();
1277
1278 if indices.null_count() == 0 {
1279 for (out_row, index) in indices.values().iter().enumerate() {
1280 let src_start = child_byte_offset + index.as_usize() * row_bytes;
1281 out_slice[out_row * row_bytes..(out_row + 1) * row_bytes]
1282 .copy_from_slice(&src[src_start..src_start + row_bytes]);
1283 }
1284 } else {
1285 for (out_row, index) in indices.values().iter().enumerate() {
1286 if indices.is_valid(out_row) {
1287 let src_start = child_byte_offset + index.as_usize() * row_bytes;
1288 out_slice[out_row * row_bytes..(out_row + 1) * row_bytes]
1289 .copy_from_slice(&src[src_start..src_start + row_bytes]);
1290 }
1291 }
1292 }
1293 let child_null_buf = taken_row_nulls.map(|n| n.expand(list_size).buffer().clone());
1295
1296 let array_data = unsafe {
1302 ArrayData::builder(child.data_type().clone())
1303 .len(out_len)
1304 .add_buffer(out.into())
1305 .null_bit_buffer(child_null_buf)
1306 .build_unchecked()
1307 };
1308 make_array(array_data)
1309}
1310
1311fn take_fixed_size_binary<IndexType: ArrowPrimitiveType, const VALIDATE_INDICES: bool>(
1317 values: &FixedSizeBinaryArray,
1318 indices: &PrimitiveArray<IndexType>,
1319 size: i32,
1320) -> Result<FixedSizeBinaryArray, ArrowError> {
1321 let size_usize = usize::try_from(size).map_err(|_| {
1322 ArrowError::InvalidArgumentError(format!("Cannot convert size '{size}' to usize"))
1323 })?;
1324
1325 let result_buffer = match size_usize {
1326 1 => take_fixed_size::<IndexType, 1>(values.values(), indices),
1327 2 => take_fixed_size::<IndexType, 2>(values.values(), indices),
1328 4 => take_fixed_size::<IndexType, 4>(values.values(), indices),
1329 8 => take_fixed_size::<IndexType, 8>(values.values(), indices),
1330 16 => take_fixed_size::<IndexType, 16>(values.values(), indices),
1331 _ => take_fixed_size_binary_buffer_dynamic_length(values, indices, size_usize),
1332 };
1333
1334 let value_nulls = take_nulls::<_, VALIDATE_INDICES>(values.nulls(), indices);
1335 let final_nulls = NullBuffer::union(value_nulls.as_ref(), indices.nulls());
1336
1337 return FixedSizeBinaryArray::try_new(size, result_buffer, final_nulls);
1338
1339 #[inline(never)]
1341 fn take_fixed_size_binary_buffer_dynamic_length<IndexType: ArrowPrimitiveType>(
1342 values: &FixedSizeBinaryArray,
1343 indices: &PrimitiveArray<IndexType>,
1344 size_usize: usize,
1345 ) -> Buffer {
1346 let values_buffer = values.values().as_slice();
1347 let mut output = Vec::with_capacity(indices.len() * size_usize);
1348
1349 if indices.null_count() == 0 {
1350 let array_iter = indices.values().iter().map(|idx| {
1351 let offset = idx.as_usize() * size_usize;
1352 &values_buffer[offset..offset + size_usize]
1353 });
1354 for slice in array_iter {
1355 output.extend_from_slice(slice);
1356 }
1357 } else {
1358 let array_iter = indices.iter().map(|idx| {
1361 idx.map(|idx| {
1362 let offset = idx.as_usize() * size_usize;
1363 &values_buffer[offset..offset + size_usize]
1364 })
1365 });
1366 for slice in array_iter {
1367 match slice {
1368 None => output.resize(output.len() + size_usize, 0),
1369 Some(slice) => output.extend_from_slice(slice),
1370 }
1371 }
1372 }
1373
1374 output.into()
1375 }
1376}
1377
1378fn take_fixed_size<IndexType: ArrowPrimitiveType, const N: usize>(
1391 buffer: &Buffer,
1392 indices: &PrimitiveArray<IndexType>,
1393) -> Buffer {
1394 assert_eq!(
1395 buffer.len() % N,
1396 0,
1397 "Invalid array length in take_fixed_size"
1398 );
1399
1400 let ptr = buffer.as_ptr();
1401 let chunk_ptr = ptr.cast::<[u8; N]>();
1402 let chunk_len = buffer.len() / N;
1403 let buffer: &[[u8; N]] = unsafe {
1404 std::slice::from_raw_parts(chunk_ptr, chunk_len)
1407 };
1408
1409 let result_buffer = match indices.nulls().filter(|n| n.null_count() > 0) {
1410 Some(n) => indices
1411 .values()
1412 .iter()
1413 .enumerate()
1414 .map(|(idx, index)| match buffer.get(index.as_usize()) {
1415 Some(v) => *v,
1416 None => match unsafe { n.inner().value_unchecked(idx) } {
1418 false => [0u8; N],
1419 true => panic!("Out-of-bounds index {index:?}"),
1420 },
1421 })
1422 .collect::<Vec<_>>(),
1423 None => indices
1424 .values()
1425 .iter()
1426 .map(|index| buffer[index.as_usize()])
1427 .collect::<Vec<_>>(),
1428 };
1429
1430 let mut vec = ManuallyDrop::new(result_buffer); let ptr = vec.as_mut_ptr();
1432 let len = vec.len();
1433 let cap = vec.capacity();
1434 let result_buffer = unsafe {
1435 Vec::from_raw_parts(ptr.cast::<u8>(), len * N, cap * N)
1437 };
1438
1439 Buffer::from_vec(result_buffer)
1440}
1441
1442fn take_dict<T: ArrowDictionaryKeyType, I: ArrowPrimitiveType, const VALIDATE_INDICES: bool>(
1447 values: &DictionaryArray<T>,
1448 indices: &PrimitiveArray<I>,
1449) -> Result<DictionaryArray<T>, ArrowError> {
1450 let new_keys = take_primitive::<_, _, VALIDATE_INDICES>(values.keys(), indices)?;
1451 Ok(unsafe { DictionaryArray::new_unchecked(new_keys, values.values().clone()) })
1452}
1453
1454fn take_run<T: RunEndIndexType, I: ArrowPrimitiveType>(
1465 run_array: &RunArray<T>,
1466 logical_indices: &PrimitiveArray<I>,
1467) -> Result<RunArray<T>, ArrowError> {
1468 if logical_indices.null_count() > 0 && !run_array.values_field().is_nullable() {
1469 return Err(ArrowError::ComputeError(
1470 "Cannot take null indices from a RunEndEncoded array with a non-nullable values field"
1471 .into(),
1472 ));
1473 }
1474
1475 let physical_indices = physical_indices_for_take(run_array, logical_indices)?;
1476
1477 let mut new_run_ends = Vec::with_capacity(1);
1481 let mut take_value_indices = Vec::with_capacity(1);
1482 let mut take_value_is_valid = NullBufferBuilder::new(1);
1483
1484 let values_cmp = make_comparator(
1485 run_array.values().as_ref(),
1486 run_array.values().as_ref(),
1487 SortOptions::default(),
1488 )?;
1489
1490 for ix in 1..physical_indices.len() {
1491 let prev_idx = physical_indices[ix - 1];
1492 let cur_idx = physical_indices[ix];
1493 if is_new_run_take(run_array.values().as_ref(), prev_idx, cur_idx, &values_cmp) {
1494 let index = I::Native::from_usize(prev_idx.unwrap_or_default()).unwrap();
1496 take_value_indices.push(index);
1497 take_value_is_valid
1498 .append(prev_idx.is_some_and(|idx| run_array.values().is_valid(idx)));
1499 new_run_ends.push(T::Native::from_usize(ix).unwrap());
1500 }
1501 }
1502 let last = physical_indices[physical_indices.len() - 1];
1503 let index = I::Native::from_usize(last.unwrap_or_default()).unwrap();
1505 take_value_indices.push(index);
1506 take_value_is_valid.append(last.is_some_and(|idx| run_array.values().is_valid(idx)));
1507 new_run_ends.push(T::Native::from_usize(physical_indices.len()).unwrap());
1508
1509 let run_ends = unsafe {
1511 RunEndBuffer::new_unchecked(ScalarBuffer::from(new_run_ends), 0, physical_indices.len())
1512 };
1513
1514 let nulls = take_value_is_valid.finish();
1515 let take_value_indices =
1516 PrimitiveArray::<I>::new(ScalarBuffer::from(take_value_indices), nulls);
1517
1518 let new_values = take(run_array.values(), &take_value_indices, None)?;
1519
1520 Ok(
1522 unsafe {
1523 RunArray::<T>::new_unchecked(run_array.data_type().clone(), run_ends, new_values)
1524 },
1525 )
1526}
1527
1528fn physical_indices_for_take<T: RunEndIndexType, I: ArrowPrimitiveType>(
1534 run_array: &RunArray<T>,
1535 logical_indices: &PrimitiveArray<I>,
1536) -> Result<Vec<Option<usize>>, ArrowError> {
1537 if logical_indices.null_count() == 0 {
1538 return Ok(run_array
1539 .get_physical_indices(logical_indices.values())?
1540 .into_iter()
1541 .map(Some)
1542 .collect());
1543 }
1544
1545 let valid_logical: Vec<_> = logical_indices.iter().flatten().collect();
1546
1547 let valid_physical = if valid_logical.is_empty() {
1548 Vec::new()
1549 } else {
1550 run_array.get_physical_indices(&valid_logical)?
1551 };
1552
1553 let mut valid_physical = valid_physical.into_iter();
1554 Ok(logical_indices
1555 .iter()
1556 .map(|index| index.map(|_| valid_physical.next().unwrap()))
1557 .collect())
1558}
1559
1560fn is_new_run_take(
1561 values: &dyn Array,
1562 prev_idx: Option<usize>,
1563 cur_idx: Option<usize>,
1564 values_cmp: &arrow_cmp::DynComparator,
1565) -> bool {
1566 let prev_valid = prev_idx.is_some_and(|idx| values.is_valid(idx));
1567 let cur_valid = cur_idx.is_some_and(|idx| values.is_valid(idx));
1568 match (prev_valid, cur_valid) {
1569 (false, false) => false,
1570 (true, true) => {
1571 let prev = prev_idx.unwrap();
1572 let cur = cur_idx.unwrap();
1573 prev != cur && values_cmp(cur, prev).is_ne()
1574 }
1575 _ => true,
1576 }
1577}
1578
1579fn take_value_indices_from_fixed_size_list<IndexType>(
1581 list: &FixedSizeListArray,
1582 indices: &PrimitiveArray<IndexType>,
1583 length: <UInt32Type as ArrowPrimitiveType>::Native,
1584) -> Result<PrimitiveArray<UInt32Type>, ArrowError>
1585where
1586 IndexType: ArrowPrimitiveType,
1587{
1588 let mut values = UInt32Builder::with_capacity(length as usize * indices.len());
1589
1590 for i in 0..indices.len() {
1591 if indices.is_valid(i) {
1592 let index = indices
1593 .value(i)
1594 .to_usize()
1595 .ok_or_else(|| ArrowError::ComputeError("Cast to usize failed".to_string()))?;
1596 let start = u32::try_from(list.value_offset_at(index)).map_err(|_| {
1597 ArrowError::ComputeError("FixedSizeList offset overflows u32".to_string())
1598 })?;
1599
1600 unsafe {
1602 values.append_trusted_len_iter(start..start + length);
1603 }
1604 } else {
1605 values.append_nulls(length as usize);
1606 }
1607 }
1608
1609 Ok(values.finish())
1610}
1611
1612trait ToIndices {
1615 type T: ArrowPrimitiveType;
1616
1617 fn to_indices(&self) -> PrimitiveArray<Self::T>;
1618}
1619
1620macro_rules! to_indices_reinterpret {
1621 ($t:ty, $o:ty) => {
1622 impl ToIndices for PrimitiveArray<$t> {
1623 type T = $o;
1624
1625 fn to_indices(&self) -> PrimitiveArray<$o> {
1626 let cast = ScalarBuffer::new(self.values().inner().clone(), 0, self.len());
1627 PrimitiveArray::new(cast, self.nulls().cloned())
1628 }
1629 }
1630 };
1631}
1632
1633macro_rules! to_indices_identity {
1634 ($t:ty) => {
1635 impl ToIndices for PrimitiveArray<$t> {
1636 type T = $t;
1637
1638 fn to_indices(&self) -> PrimitiveArray<$t> {
1639 self.clone()
1640 }
1641 }
1642 };
1643}
1644
1645macro_rules! to_indices_widening {
1646 ($t:ty, $o:ty) => {
1647 impl ToIndices for PrimitiveArray<$t> {
1648 type T = UInt32Type;
1649
1650 fn to_indices(&self) -> PrimitiveArray<$o> {
1651 let cast = self.values().iter().copied().map(|x| x as _).collect();
1652 PrimitiveArray::new(cast, self.nulls().cloned())
1653 }
1654 }
1655 };
1656}
1657
1658to_indices_widening!(UInt8Type, UInt32Type);
1659to_indices_widening!(Int8Type, UInt32Type);
1660
1661to_indices_widening!(UInt16Type, UInt32Type);
1662to_indices_widening!(Int16Type, UInt32Type);
1663
1664to_indices_identity!(UInt32Type);
1665to_indices_reinterpret!(Int32Type, UInt32Type);
1666
1667to_indices_identity!(UInt64Type);
1668to_indices_reinterpret!(Int64Type, UInt64Type);
1669
1670pub fn take_record_batch(
1709 record_batch: &RecordBatch,
1710 indices: &dyn Array,
1711) -> Result<RecordBatch, ArrowError> {
1712 downcast_integer_array!(
1713 indices => {
1714 let indices = indices.to_indices();
1715 let cols = record_batch.columns();
1716 let mut columns = Vec::with_capacity(cols.len());
1717 if let Some(first) = cols.first() {
1718 columns.push(take_impl::<_, true>(first.as_ref(), &indices)?);
1719 for col in &cols[1..] {
1720 columns.push(take_impl::<_, false>(col.as_ref(), &indices)?);
1721 }
1722 }
1723 Ok(unsafe { RecordBatch::new_unchecked(record_batch.schema(), columns, indices.len()) })
1725 },
1726 d => Err(ArrowError::InvalidArgumentError(format!("Take only supported for integers, got {d:?}")))
1727 )
1728}
1729
1730#[cfg(test)]
1731mod tests {
1732 use super::*;
1733 use arrow_array::builder::*;
1734 use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano};
1735 use arrow_schema::{Field, Fields, TimeUnit, UnionFields};
1736 use num_traits::ToPrimitive;
1737
1738 fn test_take_decimal_arrays(
1739 data: Vec<Option<i128>>,
1740 index: &UInt32Array,
1741 options: Option<TakeOptions>,
1742 expected_data: Vec<Option<i128>>,
1743 precision: &u8,
1744 scale: &i8,
1745 ) -> Result<(), ArrowError> {
1746 let output = data
1747 .into_iter()
1748 .collect::<Decimal128Array>()
1749 .with_precision_and_scale(*precision, *scale)
1750 .unwrap();
1751
1752 let expected = expected_data
1753 .into_iter()
1754 .collect::<Decimal128Array>()
1755 .with_precision_and_scale(*precision, *scale)
1756 .unwrap();
1757
1758 let expected = Arc::new(expected) as ArrayRef;
1759 let output = take(&output, index, options).unwrap();
1760 assert_eq!(&output, &expected);
1761 Ok(())
1762 }
1763
1764 fn test_take_boolean_arrays(
1765 data: Vec<Option<bool>>,
1766 index: &UInt32Array,
1767 options: Option<TakeOptions>,
1768 expected_data: Vec<Option<bool>>,
1769 ) {
1770 let output = BooleanArray::from(data);
1771 let expected = Arc::new(BooleanArray::from(expected_data)) as ArrayRef;
1772 let output = take(&output, index, options).unwrap();
1773 assert_eq!(&output, &expected)
1774 }
1775
1776 fn test_take_primitive_arrays<T>(
1777 data: Vec<Option<T::Native>>,
1778 index: &UInt32Array,
1779 options: Option<TakeOptions>,
1780 expected_data: Vec<Option<T::Native>>,
1781 ) -> Result<(), ArrowError>
1782 where
1783 T: ArrowPrimitiveType,
1784 PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
1785 {
1786 let output = PrimitiveArray::<T>::from(data);
1787 let expected = Arc::new(PrimitiveArray::<T>::from(expected_data)) as ArrayRef;
1788 let output = take(&output, index, options)?;
1789 assert_eq!(&output, &expected);
1790 Ok(())
1791 }
1792
1793 fn test_take_primitive_arrays_non_null<T>(
1794 data: Vec<T::Native>,
1795 index: &UInt32Array,
1796 options: Option<TakeOptions>,
1797 expected_data: Vec<Option<T::Native>>,
1798 ) -> Result<(), ArrowError>
1799 where
1800 T: ArrowPrimitiveType,
1801 PrimitiveArray<T>: From<Vec<T::Native>>,
1802 PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
1803 {
1804 let output = PrimitiveArray::<T>::from(data);
1805 let expected = Arc::new(PrimitiveArray::<T>::from(expected_data)) as ArrayRef;
1806 let output = take(&output, index, options)?;
1807 assert_eq!(&output, &expected);
1808 Ok(())
1809 }
1810
1811 fn test_take_impl_primitive_arrays<T, I>(
1812 data: Vec<Option<T::Native>>,
1813 index: &PrimitiveArray<I>,
1814 options: Option<TakeOptions>,
1815 expected_data: Vec<Option<T::Native>>,
1816 ) where
1817 T: ArrowPrimitiveType,
1818 PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
1819 I: ArrowPrimitiveType,
1820 {
1821 let output = PrimitiveArray::<T>::from(data);
1822 let expected = PrimitiveArray::<T>::from(expected_data);
1823 let output = take(&output, index, options).unwrap();
1824 let output = output.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();
1825 assert_eq!(output, &expected)
1826 }
1827
1828 fn create_test_struct(values: Vec<Option<(Option<bool>, Option<i32>)>>) -> StructArray {
1830 let mut struct_builder = StructBuilder::new(
1831 Fields::from(vec![
1832 Field::new("a", DataType::Boolean, true),
1833 Field::new("b", DataType::Int32, true),
1834 ]),
1835 vec![
1836 Box::new(BooleanBuilder::with_capacity(values.len())),
1837 Box::new(Int32Builder::with_capacity(values.len())),
1838 ],
1839 );
1840
1841 for value in values {
1842 struct_builder
1843 .field_builder::<BooleanBuilder>(0)
1844 .unwrap()
1845 .append_option(value.and_then(|v| v.0));
1846 struct_builder
1847 .field_builder::<Int32Builder>(1)
1848 .unwrap()
1849 .append_option(value.and_then(|v| v.1));
1850 struct_builder.append(value.is_some());
1851 }
1852 struct_builder.finish()
1853 }
1854
1855 #[test]
1856 fn test_take_decimal128_non_null_indices() {
1857 let index = UInt32Array::from(vec![0, 5, 3, 1, 4, 2]);
1858 let precision: u8 = 10;
1859 let scale: i8 = 5;
1860 test_take_decimal_arrays(
1861 vec![None, Some(3), Some(5), Some(2), Some(3), None],
1862 &index,
1863 None,
1864 vec![None, None, Some(2), Some(3), Some(3), Some(5)],
1865 &precision,
1866 &scale,
1867 )
1868 .unwrap();
1869 }
1870
1871 #[test]
1872 fn test_take_decimal128() {
1873 let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
1874 let precision: u8 = 10;
1875 let scale: i8 = 5;
1876 test_take_decimal_arrays(
1877 vec![Some(0), Some(1), Some(2), Some(3), Some(4)],
1878 &index,
1879 None,
1880 vec![Some(3), None, Some(1), Some(3), Some(2)],
1881 &precision,
1882 &scale,
1883 )
1884 .unwrap();
1885 }
1886
1887 #[test]
1888 fn test_take_primitive_non_null_indices() {
1889 let index = UInt32Array::from(vec![0, 5, 3, 1, 4, 2]);
1890 test_take_primitive_arrays::<Int8Type>(
1891 vec![None, Some(3), Some(5), Some(2), Some(3), None],
1892 &index,
1893 None,
1894 vec![None, None, Some(2), Some(3), Some(3), Some(5)],
1895 )
1896 .unwrap();
1897 }
1898
1899 #[test]
1900 fn test_take_primitive_non_null_values() {
1901 let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
1902 test_take_primitive_arrays::<Int8Type>(
1903 vec![Some(0), Some(1), Some(2), Some(3), Some(4)],
1904 &index,
1905 None,
1906 vec![Some(3), None, Some(1), Some(3), Some(2)],
1907 )
1908 .unwrap();
1909 }
1910
1911 #[test]
1912 fn test_take_primitive_non_null() {
1913 let index = UInt32Array::from(vec![0, 5, 3, 1, 4, 2]);
1914 test_take_primitive_arrays::<Int8Type>(
1915 vec![Some(0), Some(3), Some(5), Some(2), Some(3), Some(1)],
1916 &index,
1917 None,
1918 vec![Some(0), Some(1), Some(2), Some(3), Some(3), Some(5)],
1919 )
1920 .unwrap();
1921 }
1922
1923 #[test]
1924 fn test_take_primitive_nullable_indices_non_null_values_with_offset() {
1925 let index = UInt32Array::from(vec![Some(0), Some(1), Some(2), Some(3), None, None]);
1926 let index = index.slice(2, 4);
1927 let index = index.as_any().downcast_ref::<UInt32Array>().unwrap();
1928
1929 assert_eq!(
1930 index,
1931 &UInt32Array::from(vec![Some(2), Some(3), None, None])
1932 );
1933
1934 test_take_primitive_arrays_non_null::<Int64Type>(
1935 vec![0, 10, 20, 30, 40, 50],
1936 index,
1937 None,
1938 vec![Some(20), Some(30), None, None],
1939 )
1940 .unwrap();
1941 }
1942
1943 #[test]
1944 fn test_take_primitive_nullable_indices_nullable_values_with_offset() {
1945 let index = UInt32Array::from(vec![Some(0), Some(1), Some(2), Some(3), None, None]);
1946 let index = index.slice(2, 4);
1947 let index = index.as_any().downcast_ref::<UInt32Array>().unwrap();
1948
1949 assert_eq!(
1950 index,
1951 &UInt32Array::from(vec![Some(2), Some(3), None, None])
1952 );
1953
1954 test_take_primitive_arrays::<Int64Type>(
1955 vec![None, None, Some(20), Some(30), Some(40), Some(50)],
1956 index,
1957 None,
1958 vec![Some(20), Some(30), None, None],
1959 )
1960 .unwrap();
1961 }
1962
1963 #[test]
1964 fn test_take_primitive() {
1965 let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
1966
1967 test_take_primitive_arrays::<Int8Type>(
1969 vec![Some(0), None, Some(2), Some(3), None],
1970 &index,
1971 None,
1972 vec![Some(3), None, None, Some(3), Some(2)],
1973 )
1974 .unwrap();
1975
1976 test_take_primitive_arrays::<Int16Type>(
1978 vec![Some(0), None, Some(2), Some(3), None],
1979 &index,
1980 None,
1981 vec![Some(3), None, None, Some(3), Some(2)],
1982 )
1983 .unwrap();
1984
1985 test_take_primitive_arrays::<Int32Type>(
1987 vec![Some(0), None, Some(2), Some(3), None],
1988 &index,
1989 None,
1990 vec![Some(3), None, None, Some(3), Some(2)],
1991 )
1992 .unwrap();
1993
1994 test_take_primitive_arrays::<Int64Type>(
1996 vec![Some(0), None, Some(2), Some(3), None],
1997 &index,
1998 None,
1999 vec![Some(3), None, None, Some(3), Some(2)],
2000 )
2001 .unwrap();
2002
2003 test_take_primitive_arrays::<UInt8Type>(
2005 vec![Some(0), None, Some(2), Some(3), None],
2006 &index,
2007 None,
2008 vec![Some(3), None, None, Some(3), Some(2)],
2009 )
2010 .unwrap();
2011
2012 test_take_primitive_arrays::<UInt16Type>(
2014 vec![Some(0), None, Some(2), Some(3), None],
2015 &index,
2016 None,
2017 vec![Some(3), None, None, Some(3), Some(2)],
2018 )
2019 .unwrap();
2020
2021 test_take_primitive_arrays::<UInt32Type>(
2023 vec![Some(0), None, Some(2), Some(3), None],
2024 &index,
2025 None,
2026 vec![Some(3), None, None, Some(3), Some(2)],
2027 )
2028 .unwrap();
2029
2030 test_take_primitive_arrays::<Int64Type>(
2032 vec![Some(0), None, Some(2), Some(-15), None],
2033 &index,
2034 None,
2035 vec![Some(-15), None, None, Some(-15), Some(2)],
2036 )
2037 .unwrap();
2038
2039 test_take_primitive_arrays::<IntervalYearMonthType>(
2041 vec![Some(0), None, Some(2), Some(-15), None],
2042 &index,
2043 None,
2044 vec![Some(-15), None, None, Some(-15), Some(2)],
2045 )
2046 .unwrap();
2047
2048 let v1 = IntervalDayTime::new(0, 0);
2050 let v2 = IntervalDayTime::new(2, 0);
2051 let v3 = IntervalDayTime::new(-15, 0);
2052 test_take_primitive_arrays::<IntervalDayTimeType>(
2053 vec![Some(v1), None, Some(v2), Some(v3), None],
2054 &index,
2055 None,
2056 vec![Some(v3), None, None, Some(v3), Some(v2)],
2057 )
2058 .unwrap();
2059
2060 let v1 = IntervalMonthDayNano::new(0, 0, 0);
2062 let v2 = IntervalMonthDayNano::new(2, 0, 0);
2063 let v3 = IntervalMonthDayNano::new(-15, 0, 0);
2064 test_take_primitive_arrays::<IntervalMonthDayNanoType>(
2065 vec![Some(v1), None, Some(v2), Some(v3), None],
2066 &index,
2067 None,
2068 vec![Some(v3), None, None, Some(v3), Some(v2)],
2069 )
2070 .unwrap();
2071
2072 test_take_primitive_arrays::<DurationSecondType>(
2074 vec![Some(0), None, Some(2), Some(-15), None],
2075 &index,
2076 None,
2077 vec![Some(-15), None, None, Some(-15), Some(2)],
2078 )
2079 .unwrap();
2080
2081 test_take_primitive_arrays::<DurationMillisecondType>(
2083 vec![Some(0), None, Some(2), Some(-15), None],
2084 &index,
2085 None,
2086 vec![Some(-15), None, None, Some(-15), Some(2)],
2087 )
2088 .unwrap();
2089
2090 test_take_primitive_arrays::<DurationMicrosecondType>(
2092 vec![Some(0), None, Some(2), Some(-15), None],
2093 &index,
2094 None,
2095 vec![Some(-15), None, None, Some(-15), Some(2)],
2096 )
2097 .unwrap();
2098
2099 test_take_primitive_arrays::<DurationNanosecondType>(
2101 vec![Some(0), None, Some(2), Some(-15), None],
2102 &index,
2103 None,
2104 vec![Some(-15), None, None, Some(-15), Some(2)],
2105 )
2106 .unwrap();
2107
2108 test_take_primitive_arrays::<Float32Type>(
2110 vec![Some(0.0), None, Some(2.21), Some(-3.1), None],
2111 &index,
2112 None,
2113 vec![Some(-3.1), None, None, Some(-3.1), Some(2.21)],
2114 )
2115 .unwrap();
2116
2117 test_take_primitive_arrays::<Float64Type>(
2119 vec![Some(0.0), None, Some(2.21), Some(-3.1), None],
2120 &index,
2121 None,
2122 vec![Some(-3.1), None, None, Some(-3.1), Some(2.21)],
2123 )
2124 .unwrap();
2125 }
2126
2127 #[test]
2128 fn test_take_preserve_timezone() {
2129 let index = Int64Array::from(vec![Some(0), None]);
2130
2131 let input = TimestampNanosecondArray::from(vec![
2132 1_639_715_368_000_000_000,
2133 1_639_715_368_000_000_000,
2134 ])
2135 .with_timezone("UTC".to_string());
2136 let result = take(&input, &index, None).unwrap();
2137 match result.data_type() {
2138 DataType::Timestamp(TimeUnit::Nanosecond, tz) => {
2139 assert_eq!(tz.clone(), Some("UTC".into()))
2140 }
2141 _ => panic!(),
2142 }
2143 }
2144
2145 #[test]
2146 fn test_take_impl_primitive_with_int64_indices() {
2147 let index = Int64Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
2148
2149 test_take_impl_primitive_arrays::<Int16Type, Int64Type>(
2151 vec![Some(0), None, Some(2), Some(3), None],
2152 &index,
2153 None,
2154 vec![Some(3), None, None, Some(3), Some(2)],
2155 );
2156
2157 test_take_impl_primitive_arrays::<Int64Type, Int64Type>(
2159 vec![Some(0), None, Some(2), Some(-15), None],
2160 &index,
2161 None,
2162 vec![Some(-15), None, None, Some(-15), Some(2)],
2163 );
2164
2165 test_take_impl_primitive_arrays::<UInt64Type, Int64Type>(
2167 vec![Some(0), None, Some(2), Some(3), None],
2168 &index,
2169 None,
2170 vec![Some(3), None, None, Some(3), Some(2)],
2171 );
2172
2173 test_take_impl_primitive_arrays::<DurationMillisecondType, Int64Type>(
2175 vec![Some(0), None, Some(2), Some(-15), None],
2176 &index,
2177 None,
2178 vec![Some(-15), None, None, Some(-15), Some(2)],
2179 );
2180
2181 test_take_impl_primitive_arrays::<Float32Type, Int64Type>(
2183 vec![Some(0.0), None, Some(2.21), Some(-3.1), None],
2184 &index,
2185 None,
2186 vec![Some(-3.1), None, None, Some(-3.1), Some(2.21)],
2187 );
2188 }
2189
2190 #[test]
2191 fn test_take_impl_primitive_with_uint8_indices() {
2192 let index = UInt8Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
2193
2194 test_take_impl_primitive_arrays::<Int16Type, UInt8Type>(
2196 vec![Some(0), None, Some(2), Some(3), None],
2197 &index,
2198 None,
2199 vec![Some(3), None, None, Some(3), Some(2)],
2200 );
2201
2202 test_take_impl_primitive_arrays::<DurationMillisecondType, UInt8Type>(
2204 vec![Some(0), None, Some(2), Some(-15), None],
2205 &index,
2206 None,
2207 vec![Some(-15), None, None, Some(-15), Some(2)],
2208 );
2209
2210 test_take_impl_primitive_arrays::<Float32Type, UInt8Type>(
2212 vec![Some(0.0), None, Some(2.21), Some(-3.1), None],
2213 &index,
2214 None,
2215 vec![Some(-3.1), None, None, Some(-3.1), Some(2.21)],
2216 );
2217 }
2218
2219 #[test]
2220 fn test_take_bool() {
2221 let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
2222 test_take_boolean_arrays(
2224 vec![Some(false), None, Some(true), Some(false), None],
2225 &index,
2226 None,
2227 vec![Some(false), None, None, Some(false), Some(true)],
2228 );
2229 }
2230
2231 #[test]
2232 fn test_take_bool_nullable_index() {
2233 let index = UInt32Array::new(
2235 ScalarBuffer::from(vec![99, 0, 999, 1, 9999, 2]),
2236 Some(NullBuffer::from(vec![
2237 false, true, false, true, false, true,
2238 ])),
2239 );
2240 test_take_boolean_arrays(
2241 vec![Some(true), None, Some(false)],
2242 &index,
2243 None,
2244 vec![None, Some(true), None, None, None, Some(false)],
2245 );
2246 }
2247
2248 #[test]
2249 fn test_take_bool_nullable_index_nonnull_values() {
2250 let index = UInt32Array::new(
2252 ScalarBuffer::from(vec![99, 0, 999, 1, 9999, 2]),
2253 Some(NullBuffer::from(vec![
2254 false, true, false, true, false, true,
2255 ])),
2256 );
2257 test_take_boolean_arrays(
2258 vec![Some(true), Some(true), Some(false)],
2259 &index,
2260 None,
2261 vec![None, Some(true), None, Some(true), None, Some(false)],
2262 );
2263 }
2264
2265 #[test]
2266 fn test_take_bool_with_offset() {
2267 let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2), None]);
2268 let index = index.slice(2, 4);
2269 let index = index
2270 .as_any()
2271 .downcast_ref::<PrimitiveArray<UInt32Type>>()
2272 .unwrap();
2273
2274 test_take_boolean_arrays(
2276 vec![Some(false), None, Some(true), Some(false), None],
2277 index,
2278 None,
2279 vec![None, Some(false), Some(true), None],
2280 );
2281 }
2282
2283 #[test]
2284 fn test_take_bool_nullable_index_sliced_source() {
2286 let source = BooleanArray::from(vec![Some(true), Some(false), Some(true), Some(false)]);
2287 let source = source.slice(1, 3); let source = source;
2289
2290 let indices = UInt32Array::from(vec![Some(2), None, Some(0)]);
2291 let result = take(&source, &indices, None).unwrap();
2292 let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
2293
2294 let expected = BooleanArray::from(vec![Some(false), None, Some(false)]);
2295 assert_eq!(result, &expected);
2296 }
2297
2298 #[test]
2299 fn test_take_bool_no_nulls_multi_byte() {
2301 let source = BooleanArray::from(vec![
2302 true, false, true, true, false, false, true, false, true, true,
2303 ]);
2304 let indices = UInt32Array::from(vec![0, 2, 4, 6, 8, 1, 3, 5, 7, 9]);
2305 let result = take(&source, &indices, None).unwrap();
2306 let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
2307 let expected = BooleanArray::from(vec![
2308 true, true, false, true, true, false, true, false, false, true,
2309 ]);
2310 assert_eq!(result, &expected);
2311 }
2312
2313 #[test]
2314 fn test_take_bool_no_nulls_sliced_source() {
2316 let source = BooleanArray::from(vec![true, false, true, false, true]);
2317 let source = source.slice(2, 3); let source = source.as_any().downcast_ref::<BooleanArray>().unwrap();
2319 let indices = UInt32Array::from(vec![2, 0, 1]);
2320 let result = take(source, &indices, None).unwrap();
2321 let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
2322 let expected = BooleanArray::from(vec![true, true, false]);
2323 assert_eq!(result, &expected);
2324 }
2325
2326 #[test]
2327 fn test_take_bool_nullable_values_multi_byte() {
2329 let source = BooleanArray::from(vec![
2330 Some(true),
2331 None,
2332 Some(false),
2333 Some(true),
2334 None,
2335 Some(false),
2336 Some(true),
2337 Some(false),
2338 Some(true),
2339 None,
2340 ]);
2341 let indices = UInt32Array::from(vec![0, 2, 4, 6, 8, 1, 3, 5, 7, 9]);
2342 let result = take(&source, &indices, None).unwrap();
2343 let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
2344 let expected = BooleanArray::from(vec![
2345 Some(true),
2346 Some(false),
2347 None,
2348 Some(true),
2349 Some(true),
2350 None,
2351 Some(true),
2352 Some(false),
2353 Some(false),
2354 None,
2355 ]);
2356 assert_eq!(result, &expected);
2357 }
2358
2359 #[test]
2360 fn test_take_bool_nullable_values_sliced_source_null_indices() {
2362 let source =
2363 BooleanArray::from(vec![Some(true), Some(false), None, Some(true), Some(false)]);
2364 let source = source.slice(1, 4); let source = source.as_any().downcast_ref::<BooleanArray>().unwrap();
2366 let indices = UInt32Array::from(vec![Some(3), None, Some(1), Some(0)]);
2367 let result = take(source, &indices, None).unwrap();
2368 let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
2369 let expected = BooleanArray::from(vec![Some(false), None, None, Some(false)]);
2370 assert_eq!(result, &expected);
2371 }
2372
2373 #[test]
2374 #[should_panic(expected = "assertion failed: idx < self.bit_len")]
2375 fn test_take_bool_oob_no_check_bounds_panics() {
2376 let array = BooleanArray::from(vec![true, false, true]);
2377 let indices = Int32Array::from(vec![0, 1, 10]);
2378 take(&array, &indices, None).unwrap();
2379 }
2380
2381 fn _test_take_string<'a, K>()
2382 where
2383 K: Array + PartialEq + From<Vec<Option<&'a str>>> + 'static,
2384 {
2385 let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(4)]);
2386
2387 let array = K::from(vec![
2388 Some("one"),
2389 None,
2390 Some("three"),
2391 Some("four"),
2392 Some("five"),
2393 ]);
2394 let actual = take(&array, &index, None).unwrap();
2395 assert_eq!(actual.len(), index.len());
2396
2397 let actual = actual.as_any().downcast_ref::<K>().unwrap();
2398
2399 let expected = K::from(vec![Some("four"), None, None, Some("four"), Some("five")]);
2400
2401 assert_eq!(actual, &expected);
2402 }
2403
2404 #[test]
2405 fn test_take_string() {
2406 _test_take_string::<StringArray>()
2407 }
2408
2409 #[test]
2410 fn test_take_large_string() {
2411 _test_take_string::<LargeStringArray>()
2412 }
2413
2414 #[test]
2415 fn test_take_slice_string() {
2416 let strings = StringArray::from(vec![Some("hello"), None, Some("world"), None, Some("hi")]);
2417 let indices = Int32Array::from(vec![Some(0), Some(1), None, Some(0), Some(2)]);
2418 let indices_slice = indices.slice(1, 4);
2419 let expected = StringArray::from(vec![None, None, Some("hello"), Some("world")]);
2420 let result = take(&strings, &indices_slice, None).unwrap();
2421 assert_eq!(result.as_ref(), &expected);
2422 }
2423
2424 #[test]
2429 fn test_take_bytes_sliced_values() {
2430 let values = StringArray::from(vec![
2431 Some("aaa"),
2432 Some("bbb"),
2433 None,
2434 Some("ccccc"),
2435 Some("dd"),
2436 None,
2437 Some("eeee"),
2438 ]);
2439 let sliced = values.slice(2, 5);
2442
2443 let indices = Int32Array::from(vec![1, 2, 4, 1]);
2446 let result = take(&sliced, &indices, None).unwrap();
2447 let expected =
2448 StringArray::from(vec![Some("ccccc"), Some("dd"), Some("eeee"), Some("ccccc")]);
2449 assert_eq!(result.as_string::<i32>(), &expected);
2450
2451 let indices = Int32Array::from(vec![Some(1), None, Some(0), Some(4), Some(3)]);
2454 let result = take(&sliced, &indices, None).unwrap();
2455 let expected = StringArray::from(vec![Some("ccccc"), None, None, Some("eeee"), None]);
2456 assert_eq!(result.as_string::<i32>(), &expected);
2457 }
2458
2459 fn _test_byte_view<T>()
2460 where
2461 T: ByteViewType,
2462 str: AsRef<T::Native>,
2463 T::Native: PartialEq,
2464 {
2465 let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(4), Some(2)]);
2466 let array = {
2467 let mut builder = GenericByteViewBuilder::<T>::new();
2469 builder.append_value("hello");
2470 builder.append_value("world");
2471 builder.append_null();
2472 builder.append_value("large payload over 12 bytes");
2473 builder.append_value("lulu");
2474 builder.finish()
2475 };
2476
2477 let actual = take(&array, &index, None).unwrap();
2478
2479 assert_eq!(actual.len(), index.len());
2480 let actual_buffers = actual.as_byte_view::<T>().data_buffers();
2481 let input_buffers = array.data_buffers();
2482 assert!(Arc::ptr_eq(actual_buffers, input_buffers));
2483
2484 let expected = {
2485 let mut builder = GenericByteViewBuilder::<T>::new();
2487 builder.append_value("large payload over 12 bytes");
2488 builder.append_null();
2489 builder.append_value("world");
2490 builder.append_value("large payload over 12 bytes");
2491 builder.append_value("lulu");
2492 builder.append_null();
2493 builder.finish()
2494 };
2495
2496 assert_eq!(actual.as_ref(), &expected);
2497 }
2498
2499 #[test]
2500 fn test_take_string_view() {
2501 _test_byte_view::<StringViewType>()
2502 }
2503
2504 #[test]
2505 fn test_take_binary_view() {
2506 _test_byte_view::<BinaryViewType>()
2507 }
2508
2509 macro_rules! test_take_list {
2510 ($offset_type:ty, $list_data_type:ident, $list_array_type:ident) => {{
2511 let values = Arc::new(Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2, 3]));
2513 let value_offsets =
2514 OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0, 3, 6, 6, 8]));
2515 let field = Arc::new(Field::new_list_field(DataType::Int32, false));
2516 let list_array = $list_array_type::new(Arc::clone(&field), value_offsets, values, None);
2517
2518 let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(2), Some(0)]);
2520
2521 let a = take(&list_array, &index, None).unwrap();
2522 let a: &$list_array_type = a.as_any().downcast_ref::<$list_array_type>().unwrap();
2523
2524 let expected_values = Arc::new(Int32Array::from(vec![
2527 Some(2),
2528 Some(3),
2529 Some(-1),
2530 Some(-2),
2531 Some(-1),
2532 Some(0),
2533 Some(0),
2534 Some(0),
2535 ]));
2536 let expected_offsets =
2537 OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0, 2, 2, 5, 5, 8]));
2538 let expected_list_array = $list_array_type::new(
2539 field,
2540 expected_offsets,
2541 expected_values,
2542 index.nulls().cloned(),
2544 );
2545
2546 assert_eq!(a, &expected_list_array);
2547 }};
2548 }
2549
2550 macro_rules! test_take_list_with_value_nulls {
2551 ($offset_type:ty, $list_data_type:ident, $list_array_type:ident) => {{
2552 let values = Arc::new(Int32Array::from(vec![
2554 Some(0),
2555 None,
2556 Some(0),
2557 Some(-1),
2558 Some(-2),
2559 Some(3),
2560 None,
2561 Some(5),
2562 None,
2563 ]));
2564 let value_offsets =
2565 OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0, 3, 6, 7, 9]));
2566 let field = Arc::new(Field::new_list_field(DataType::Int32, true));
2567 let list_array = $list_array_type::new(Arc::clone(&field), value_offsets, values, None);
2568
2569 let index = UInt32Array::from(vec![Some(2), None, Some(1), Some(3), Some(0)]);
2571
2572 let a = take(&list_array, &index, None).unwrap();
2573 let a: &$list_array_type = a.as_any().downcast_ref::<$list_array_type>().unwrap();
2574
2575 let expected_values = Arc::new(Int32Array::from(vec![
2578 None,
2579 Some(-1),
2580 Some(-2),
2581 Some(3),
2582 Some(5),
2583 None,
2584 Some(0),
2585 None,
2586 Some(0),
2587 ]));
2588 let expected_offsets =
2589 OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0, 1, 1, 4, 6, 9]));
2590 let expected_list_array = $list_array_type::new(
2591 field,
2592 expected_offsets,
2593 expected_values,
2594 index.nulls().cloned(),
2596 );
2597
2598 assert_eq!(a, &expected_list_array);
2599 }};
2600 }
2601
2602 macro_rules! test_take_list_with_nulls {
2603 ($offset_type:ty, $list_data_type:ident, $list_array_type:ident) => {{
2604 let values = Arc::new(Int32Array::from(vec![
2606 Some(0),
2607 None,
2608 Some(0),
2609 Some(-1),
2610 Some(-2),
2611 Some(3),
2612 Some(5),
2613 None,
2614 ]));
2615 let value_offsets =
2616 OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0, 3, 6, 6, 8]));
2617 let field = Arc::new(Field::new_list_field(DataType::Int32, true));
2618 let list_nulls = NullBuffer::from(vec![true, true, false, true]);
2620 let list_array =
2621 $list_array_type::new(Arc::clone(&field), value_offsets, values, Some(list_nulls));
2622
2623 let index = UInt32Array::from(vec![Some(2), None, Some(1), Some(3), Some(0)]);
2625
2626 let a = take(&list_array, &index, None).unwrap();
2627 let a: &$list_array_type = a.as_any().downcast_ref::<$list_array_type>().unwrap();
2628
2629 let expected_values = Arc::new(Int32Array::from(vec![
2632 Some(-1),
2633 Some(-2),
2634 Some(3),
2635 Some(5),
2636 None,
2637 Some(0),
2638 None,
2639 Some(0),
2640 ]));
2641 let expected_offsets =
2642 OffsetBuffer::<$offset_type>::new(ScalarBuffer::from(vec![0, 0, 0, 3, 5, 8]));
2643 let expected_nulls = NullBuffer::from(vec![false, false, true, true, true]);
2645 let expected_list_array = $list_array_type::new(
2646 field,
2647 expected_offsets,
2648 expected_values,
2649 Some(expected_nulls),
2650 );
2651
2652 assert_eq!(a, &expected_list_array);
2653 }};
2654 }
2655
2656 fn test_take_list_view_generic<OffsetType: OffsetSizeTrait, ValuesType: ArrowPrimitiveType, F>(
2657 values: Vec<Option<Vec<Option<ValuesType::Native>>>>,
2658 take_indices: Vec<Option<usize>>,
2659 expected: Vec<Option<Vec<Option<ValuesType::Native>>>>,
2660 mapper: F,
2661 ) where
2662 F: Fn(GenericListViewArray<OffsetType>) -> GenericListViewArray<OffsetType>,
2663 {
2664 let mut list_view_array =
2665 GenericListViewBuilder::<OffsetType, _>::new(PrimitiveBuilder::<ValuesType>::new());
2666
2667 for value in values {
2668 list_view_array.append_option(value);
2669 }
2670 let list_view_array = list_view_array.finish();
2671 let list_view_array = mapper(list_view_array);
2672
2673 let mut indices = UInt64Builder::new();
2674 for idx in take_indices {
2675 indices.append_option(idx.map(|i| i.to_u64().unwrap()));
2676 }
2677 let indices = indices.finish();
2678
2679 let taken = take(&list_view_array, &indices, None)
2680 .unwrap()
2681 .as_list_view()
2682 .clone();
2683
2684 let mut expected_array =
2685 GenericListViewBuilder::<OffsetType, _>::new(PrimitiveBuilder::<ValuesType>::new());
2686 for value in expected {
2687 expected_array.append_option(value);
2688 }
2689 let expected_array = expected_array.finish();
2690
2691 assert_eq!(taken, expected_array);
2692 }
2693
2694 macro_rules! list_view_test_case {
2695 (values: $values:expr, indices: $indices:expr, expected: $expected: expr) => {{
2696 test_take_list_view_generic::<i32, Int8Type, _>($values, $indices, $expected, |x| x);
2697 test_take_list_view_generic::<i64, Int8Type, _>($values, $indices, $expected, |x| x);
2698 }};
2699 (values: $values:expr, transform: $fn:expr, indices: $indices:expr, expected: $expected: expr) => {{
2700 test_take_list_view_generic::<i32, Int8Type, _>($values, $indices, $expected, $fn);
2701 test_take_list_view_generic::<i64, Int8Type, _>($values, $indices, $expected, $fn);
2702 }};
2703 }
2704
2705 fn do_take_fixed_size_list_test<T>(
2706 length: <Int32Type as ArrowPrimitiveType>::Native,
2707 input_data: Vec<Option<Vec<Option<T::Native>>>>,
2708 indices: Vec<<UInt32Type as ArrowPrimitiveType>::Native>,
2709 expected_data: Vec<Option<Vec<Option<T::Native>>>>,
2710 ) where
2711 T: ArrowPrimitiveType,
2712 PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
2713 {
2714 let indices = UInt32Array::from(indices);
2715
2716 let input_array = FixedSizeListArray::from_iter_primitive::<T, _, _>(input_data, length);
2717
2718 let output =
2719 take_fixed_size_list::<_, true>(&input_array, &indices, length as u32).unwrap();
2720
2721 let expected = FixedSizeListArray::from_iter_primitive::<T, _, _>(expected_data, length);
2722
2723 assert_eq!(&output, &expected)
2724 }
2725
2726 #[test]
2727 fn test_take_list_primitive_child_null_indices() {
2729 let list = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
2732 Some(vec![Some(1)]),
2733 Some(vec![Some(2), Some(3), Some(4)]),
2734 Some(vec![Some(5), Some(6)]),
2735 ]);
2736 let indices = Int32Array::from(vec![Some(2), None, Some(0), Some(1)]);
2737 let result = take(&list, &indices, None).unwrap();
2738 let result = result.as_any().downcast_ref::<ListArray>().unwrap();
2739
2740 let expected = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
2741 Some(vec![Some(5), Some(6)]),
2742 None,
2743 Some(vec![Some(1)]),
2744 Some(vec![Some(2), Some(3), Some(4)]),
2745 ]);
2746 assert_eq!(result, &expected);
2747 }
2748
2749 #[test]
2750 fn test_take_list() {
2751 test_take_list!(i32, List, ListArray);
2752 }
2753
2754 #[test]
2755 fn test_take_large_list() {
2756 test_take_list!(i64, LargeList, LargeListArray);
2757 }
2758
2759 #[test]
2760 fn test_take_list_with_value_nulls() {
2761 test_take_list_with_value_nulls!(i32, List, ListArray);
2762 }
2763
2764 #[test]
2765 fn test_take_large_list_with_value_nulls() {
2766 test_take_list_with_value_nulls!(i64, LargeList, LargeListArray);
2767 }
2768
2769 #[test]
2770 fn test_test_take_list_with_nulls() {
2771 test_take_list_with_nulls!(i32, List, ListArray);
2772 }
2773
2774 #[test]
2775 fn test_test_take_large_list_with_nulls() {
2776 test_take_list_with_nulls!(i64, LargeList, LargeListArray);
2777 }
2778
2779 #[test]
2780 fn test_test_take_list_view_reversed() {
2781 list_view_test_case! {
2783 values: vec![
2784 Some(vec![Some(1), None, Some(3)]),
2785 None,
2786 Some(vec![Some(7), Some(8), None]),
2787 ],
2788 indices: vec![Some(2), Some(1), Some(0)],
2789 expected: vec![
2790 Some(vec![Some(7), Some(8), None]),
2791 None,
2792 Some(vec![Some(1), None, Some(3)]),
2793 ]
2794 }
2795 }
2796
2797 #[test]
2798 fn test_take_list_view_null_indices() {
2799 list_view_test_case! {
2801 values: vec![
2802 Some(vec![Some(1), None, Some(3)]),
2803 None,
2804 Some(vec![Some(7), Some(8), None]),
2805 ],
2806 indices: vec![None, Some(0), None],
2807 expected: vec![None, Some(vec![Some(1), None, Some(3)]), None]
2808 }
2809 }
2810
2811 #[test]
2812 fn test_take_list_view_null_values() {
2813 list_view_test_case! {
2815 values: vec![
2816 Some(vec![Some(1), None, Some(3)]),
2817 None,
2818 Some(vec![Some(7), Some(8), None]),
2819 ],
2820 indices: vec![Some(1), Some(1), Some(1), None, None],
2821 expected: vec![None; 5]
2822 }
2823 }
2824
2825 #[test]
2826 fn test_take_list_view_sliced() {
2827 list_view_test_case! {
2829 values: vec![
2830 Some(vec![Some(1)]),
2831 None,
2832 None,
2833 Some(vec![Some(2), Some(3)]),
2834 Some(vec![Some(4), Some(5)]),
2835 None,
2836 ],
2837 transform: |l| l.slice(2, 4),
2838 indices: vec![Some(0), Some(3), None, Some(1), Some(2)],
2839 expected: vec![
2840 None, None, None, Some(vec![Some(2), Some(3)]), Some(vec![Some(4), Some(5)])
2841 ]
2842 }
2843 }
2844
2845 #[test]
2846 fn test_take_fixed_size_list() {
2847 do_take_fixed_size_list_test::<Int32Type>(
2848 3,
2849 vec![
2850 Some(vec![None, Some(1), Some(2)]),
2851 Some(vec![Some(3), Some(4), None]),
2852 Some(vec![Some(6), Some(7), Some(8)]),
2853 ],
2854 vec![2, 1, 0],
2855 vec![
2856 Some(vec![Some(6), Some(7), Some(8)]),
2857 Some(vec![Some(3), Some(4), None]),
2858 Some(vec![None, Some(1), Some(2)]),
2859 ],
2860 );
2861
2862 do_take_fixed_size_list_test::<UInt8Type>(
2863 1,
2864 vec![
2865 Some(vec![Some(1)]),
2866 Some(vec![Some(2)]),
2867 Some(vec![Some(3)]),
2868 Some(vec![Some(4)]),
2869 Some(vec![Some(5)]),
2870 Some(vec![Some(6)]),
2871 Some(vec![Some(7)]),
2872 Some(vec![Some(8)]),
2873 ],
2874 vec![2, 7, 0],
2875 vec![
2876 Some(vec![Some(3)]),
2877 Some(vec![Some(8)]),
2878 Some(vec![Some(1)]),
2879 ],
2880 );
2881
2882 do_take_fixed_size_list_test::<UInt64Type>(
2883 3,
2884 vec![
2885 Some(vec![Some(10), Some(11), Some(12)]),
2886 Some(vec![Some(13), Some(14), Some(15)]),
2887 None,
2888 Some(vec![Some(16), Some(17), Some(18)]),
2889 ],
2890 vec![3, 2, 1, 2, 0],
2891 vec![
2892 Some(vec![Some(16), Some(17), Some(18)]),
2893 None,
2894 Some(vec![Some(13), Some(14), Some(15)]),
2895 None,
2896 Some(vec![Some(10), Some(11), Some(12)]),
2897 ],
2898 );
2899 }
2900
2901 #[test]
2902 fn test_take_fixed_size_binary_with_nulls_indices() {
2903 let fsb = FixedSizeBinaryArray::try_from_sparse_iter_with_size(
2904 [
2905 Some(vec![0x01, 0x01, 0x01, 0x01]),
2906 Some(vec![0x02, 0x02, 0x02, 0x02]),
2907 Some(vec![0x03, 0x03, 0x03, 0x03]),
2908 Some(vec![0x04, 0x04, 0x04, 0x04]),
2909 ]
2910 .into_iter(),
2911 4,
2912 )
2913 .unwrap();
2914
2915 let indices = UInt32Array::from(vec![Some(0), None, None, Some(3)]);
2917
2918 let result = take_fixed_size_binary::<_, true>(&fsb, &indices, 4).unwrap();
2919 assert_eq!(result.len(), 4);
2920 assert_eq!(result.null_count(), 2);
2921 assert_eq!(
2922 result.nulls().unwrap().iter().collect::<Vec<_>>(),
2923 vec![true, false, false, true]
2924 );
2925 }
2926
2927 #[test]
2931 fn test_take_fixed_size_binary_with_nulls_indices_not_optimized_length() {
2932 let fsb = FixedSizeBinaryArray::try_from_sparse_iter_with_size(
2933 [
2934 Some(vec![0x01, 0x01, 0x01, 0x01, 0x01]),
2935 Some(vec![0x02, 0x02, 0x02, 0x02, 0x01]),
2936 Some(vec![0x03, 0x03, 0x03, 0x03, 0x01]),
2937 Some(vec![0x04, 0x04, 0x04, 0x04, 0x01]),
2938 ]
2939 .into_iter(),
2940 5,
2941 )
2942 .unwrap();
2943
2944 let indices = UInt32Array::from(vec![Some(0), None, None, Some(3)]);
2946
2947 let result = take_fixed_size_binary::<_, true>(&fsb, &indices, 5).unwrap();
2948 assert_eq!(result.len(), 4);
2949 assert_eq!(result.null_count(), 2);
2950 assert_eq!(
2951 result.nulls().unwrap().iter().collect::<Vec<_>>(),
2952 vec![true, false, false, true]
2953 );
2954 }
2955
2956 #[test]
2957 #[should_panic(expected = "index out of bounds: the len is 4 but the index is 1000")]
2958 fn test_take_list_out_of_bounds() {
2959 let values = Arc::new(Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2, 3]));
2961 let value_offsets = OffsetBuffer::<i32>::new(ScalarBuffer::from(vec![0, 3, 6, 8]));
2963 let field = Arc::new(Field::new_list_field(DataType::Int32, false));
2965 let list_array = ListArray::new(field, value_offsets, values, None);
2966
2967 let index = UInt32Array::from(vec![1000]);
2968
2969 take(&list_array, &index, None).unwrap();
2972 }
2973
2974 #[test]
2975 fn test_take_map() {
2976 let values = Int32Array::from(vec![1, 2, 3, 4]);
2977 let array =
2978 MapArray::new_from_strings(vec!["a", "b", "c", "a"].into_iter(), &values, &[0, 3, 4])
2979 .unwrap();
2980
2981 let index = UInt32Array::from(vec![0]);
2982
2983 let result = take(&array, &index, None).unwrap();
2984 let expected: ArrayRef = Arc::new(
2985 MapArray::new_from_strings(
2986 vec!["a", "b", "c"].into_iter(),
2987 &values.slice(0, 3),
2988 &[0, 3],
2989 )
2990 .unwrap(),
2991 );
2992 assert_eq!(&expected, &result);
2993 }
2994
2995 #[test]
2996 fn test_take_struct() {
2997 let array = create_test_struct(vec![
2998 Some((Some(true), Some(42))),
2999 Some((Some(false), Some(28))),
3000 Some((Some(false), Some(19))),
3001 Some((Some(true), Some(31))),
3002 None,
3003 ]);
3004
3005 let index = UInt32Array::from(vec![0, 3, 1, 0, 2, 4]);
3006 let actual = take(&array, &index, None).unwrap();
3007 let actual: &StructArray = actual.as_any().downcast_ref::<StructArray>().unwrap();
3008 assert_eq!(index.len(), actual.len());
3009 assert_eq!(1, actual.null_count());
3010
3011 let expected = create_test_struct(vec![
3012 Some((Some(true), Some(42))),
3013 Some((Some(true), Some(31))),
3014 Some((Some(false), Some(28))),
3015 Some((Some(true), Some(42))),
3016 Some((Some(false), Some(19))),
3017 None,
3018 ]);
3019
3020 assert_eq!(&expected, actual);
3021
3022 let nulls = NullBuffer::from(&[false, true, false, true, false, true]);
3023 let empty_struct_arr = StructArray::new_empty_fields(6, Some(nulls));
3024 let index = UInt32Array::from(vec![0, 2, 1, 4]);
3025 let actual = take(&empty_struct_arr, &index, None).unwrap();
3026
3027 let expected_nulls = NullBuffer::from(&[false, false, true, false]);
3028 let expected_struct_arr = StructArray::new_empty_fields(4, Some(expected_nulls));
3029 assert_eq!(&expected_struct_arr, actual.as_struct());
3030 }
3031
3032 #[test]
3033 fn test_take_struct_with_null_indices() {
3034 let array = create_test_struct(vec![
3035 Some((Some(true), Some(42))),
3036 Some((Some(false), Some(28))),
3037 Some((Some(false), Some(19))),
3038 Some((Some(true), Some(31))),
3039 None,
3040 ]);
3041
3042 let index = UInt32Array::from(vec![None, Some(3), Some(1), None, Some(0), Some(4)]);
3043 let actual = take(&array, &index, None).unwrap();
3044 let actual: &StructArray = actual.as_any().downcast_ref::<StructArray>().unwrap();
3045 assert_eq!(index.len(), actual.len());
3046 assert_eq!(3, actual.null_count()); let expected = create_test_struct(vec![
3049 None,
3050 Some((Some(true), Some(31))),
3051 Some((Some(false), Some(28))),
3052 None,
3053 Some((Some(true), Some(42))),
3054 None,
3055 ]);
3056
3057 assert_eq!(&expected, actual);
3058 }
3059
3060 #[test]
3061 fn test_take_out_of_bounds() {
3062 let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(6)]);
3063 let take_opt = TakeOptions { check_bounds: true };
3064
3065 let result = test_take_primitive_arrays::<Int64Type>(
3067 vec![Some(0), None, Some(2), Some(3), None],
3068 &index,
3069 Some(take_opt),
3070 vec![None],
3071 );
3072 assert!(result.is_err());
3073 }
3074
3075 #[test]
3076 #[should_panic(expected = "index out of bounds: the len is 4 but the index is 1000")]
3077 fn test_take_out_of_bounds_panic() {
3078 let index = UInt32Array::from(vec![Some(1000)]);
3079
3080 test_take_primitive_arrays::<Int64Type>(
3081 vec![Some(0), Some(1), Some(2), Some(3)],
3082 &index,
3083 None,
3084 vec![None],
3085 )
3086 .unwrap();
3087 }
3088
3089 #[test]
3090 fn test_null_array_smaller_than_indices() {
3091 let values = NullArray::new(2);
3092 let indices = UInt32Array::from(vec![Some(0), None, Some(15)]);
3093
3094 let result = take(&values, &indices, None).unwrap();
3095 let expected: ArrayRef = Arc::new(NullArray::new(3));
3096 assert_eq!(&result, &expected);
3097 }
3098
3099 #[test]
3100 fn test_null_array_larger_than_indices() {
3101 let values = NullArray::new(5);
3102 let indices = UInt32Array::from(vec![Some(0), None, Some(15)]);
3103
3104 let result = take(&values, &indices, None).unwrap();
3105 let expected: ArrayRef = Arc::new(NullArray::new(3));
3106 assert_eq!(&result, &expected);
3107 }
3108
3109 #[test]
3110 fn test_null_array_indices_out_of_bounds() {
3111 let values = NullArray::new(5);
3112 let indices = UInt32Array::from(vec![Some(0), None, Some(15)]);
3113
3114 let result = take(&values, &indices, Some(TakeOptions { check_bounds: true }));
3115 assert_eq!(
3116 result.unwrap_err().to_string(),
3117 "Compute error: Array index out of bounds, cannot get item at index 15 from 5 entries"
3118 );
3119 }
3120
3121 #[test]
3122 fn test_take_dict() {
3123 let mut dict_builder = StringDictionaryBuilder::<Int16Type>::new();
3124
3125 dict_builder.append("foo").unwrap();
3126 dict_builder.append("bar").unwrap();
3127 dict_builder.append("").unwrap();
3128 dict_builder.append_null();
3129 dict_builder.append("foo").unwrap();
3130 dict_builder.append("bar").unwrap();
3131 dict_builder.append("bar").unwrap();
3132 dict_builder.append("foo").unwrap();
3133
3134 let array = dict_builder.finish();
3135 let dict_values = array.values().clone();
3136 let dict_values = dict_values.as_any().downcast_ref::<StringArray>().unwrap();
3137
3138 let indices = UInt32Array::from(vec![
3139 Some(0), Some(7), None, Some(5), Some(6), Some(2), Some(3), ]);
3147
3148 let result = take(&array, &indices, None).unwrap();
3149 let result = result
3150 .as_any()
3151 .downcast_ref::<DictionaryArray<Int16Type>>()
3152 .unwrap();
3153
3154 let result_values: StringArray = result.values().to_data().into();
3155
3156 let expected_values = StringArray::from(vec!["foo", "bar", ""]);
3158 assert_eq!(&expected_values, dict_values);
3159 assert_eq!(&expected_values, &result_values);
3160
3161 let expected_keys = Int16Array::from(vec![
3162 Some(0),
3163 Some(0),
3164 None,
3165 Some(1),
3166 Some(1),
3167 Some(2),
3168 None,
3169 ]);
3170 assert_eq!(result.keys(), &expected_keys);
3171 }
3172
3173 fn build_generic_list<S, T>(data: Vec<Option<Vec<T::Native>>>) -> GenericListArray<S>
3174 where
3175 S: OffsetSizeTrait + 'static,
3176 T: ArrowPrimitiveType,
3177 PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
3178 {
3179 GenericListArray::from_iter_primitive::<T, _, _>(
3180 data.iter()
3181 .map(|x| x.as_ref().map(|x| x.iter().map(|x| Some(*x)))),
3182 )
3183 }
3184
3185 fn test_take_sliced_list_generic<S: OffsetSizeTrait + 'static>() {
3186 let list = build_generic_list::<S, Int32Type>(vec![
3187 Some(vec![0, 1]),
3188 Some(vec![2, 3, 4]),
3189 None,
3190 Some(vec![]),
3191 Some(vec![5, 6]),
3192 Some(vec![7]),
3193 ]);
3194 let sliced = list.slice(1, 4);
3195 let indices = UInt32Array::from(vec![Some(3), Some(0), None, Some(2), Some(1)]);
3196
3197 let taken = take(&sliced, &indices, None).unwrap();
3198 let taken = taken.as_list::<S>();
3199
3200 let expected = build_generic_list::<S, Int32Type>(vec![
3201 Some(vec![5, 6]),
3202 Some(vec![2, 3, 4]),
3203 None,
3204 Some(vec![]),
3205 None,
3206 ]);
3207
3208 assert_eq!(taken, &expected);
3209 }
3210
3211 fn test_take_sliced_list_with_value_nulls_generic<S: OffsetSizeTrait + 'static>() {
3212 let list = GenericListArray::<S>::from_iter_primitive::<Int32Type, _, _>(vec![
3213 Some(vec![Some(10)]),
3214 Some(vec![None, Some(1)]),
3215 None,
3216 Some(vec![Some(2), None]),
3217 Some(vec![]),
3218 Some(vec![Some(3)]),
3219 ]);
3220 let sliced = list.slice(1, 4);
3221 let indices = UInt32Array::from(vec![Some(2), Some(0), None, Some(3), Some(1)]);
3222
3223 let taken = take(&sliced, &indices, None).unwrap();
3224 let taken = taken.as_list::<S>();
3225
3226 let expected = GenericListArray::<S>::from_iter_primitive::<Int32Type, _, _>(vec![
3227 Some(vec![Some(2), None]),
3228 Some(vec![None, Some(1)]),
3229 None,
3230 Some(vec![]),
3231 None,
3232 ]);
3233
3234 assert_eq!(taken, &expected);
3235 }
3236
3237 #[test]
3238 fn test_take_sliced_list() {
3239 test_take_sliced_list_generic::<i32>();
3240 }
3241
3242 #[test]
3243 fn test_take_sliced_large_list() {
3244 test_take_sliced_list_generic::<i64>();
3245 }
3246
3247 #[test]
3248 fn test_take_sliced_list_with_value_nulls() {
3249 test_take_sliced_list_with_value_nulls_generic::<i32>();
3250 }
3251
3252 #[test]
3253 fn test_take_sliced_large_list_with_value_nulls() {
3254 test_take_sliced_list_with_value_nulls_generic::<i64>();
3255 }
3256
3257 #[test]
3258 fn test_take_sliced_list_memory_size() {
3259 fn check(parent: ArrayRef, unsliced: ArrayRef) {
3262 let sliced = parent.slice(500, 2);
3263 let indices = UInt32Array::from_iter_values((0..100).map(|i| i % 2));
3264 let actual = take(&sliced, &indices, None).unwrap();
3265 let expected = take(&unsliced, &indices, None).unwrap();
3266 assert_eq!(&actual, &expected);
3267 assert_eq!(
3268 actual.get_array_memory_size(),
3269 expected.get_array_memory_size()
3270 );
3271 }
3272
3273 let list = |rows: usize, values: ArrayRef| -> ArrayRef {
3275 let field = Arc::new(Field::new_list_field(values.data_type().clone(), false));
3276 let offsets = OffsetBuffer::from_repeated_length(10, rows);
3277 Arc::new(ListArray::new(field, offsets, values, None))
3278 };
3279 check(
3280 list(1000, Arc::new(Int64Array::from(vec![1; 10_000]))),
3281 list(2, Arc::new(Int64Array::from(vec![1; 20]))),
3282 );
3283 check(
3284 list(1000, Arc::new(StringArray::from(vec!["a"; 10_000]))),
3285 list(2, Arc::new(StringArray::from(vec!["a"; 20]))),
3286 );
3287 }
3288
3289 #[test]
3290 fn test_take_runs() {
3291 let logical_array: Vec<i32> = vec![1_i32, 1, 2, 2, 1, 1, 1, 2, 2, 1, 1, 2, 2];
3292
3293 let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
3294 builder.extend(logical_array.into_iter().map(Some));
3295 let run_array = builder.finish();
3296
3297 let take_indices: PrimitiveArray<Int32Type> =
3298 vec![7, 2, 3, 7, 11, 4, 6].into_iter().collect();
3299
3300 let take_out = take_run(&run_array, &take_indices).unwrap();
3301
3302 assert_eq!(take_out.len(), 7);
3303 assert_eq!(
3305 take_out.run_ends().values().len(),
3306 2,
3307 "expected two physical runs"
3308 );
3309 assert_eq!(take_out.run_ends().values(), &[5_i32, 7]);
3310
3311 let take_out_values = take_out.values().as_primitive::<Int32Type>();
3312 assert_eq!(take_out_values.values(), &[2, 1]);
3313 }
3314
3315 #[test]
3316 fn test_take_runs_null_indices() {
3317 let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
3320 builder.extend([Some(10), Some(10), None, None, Some(99)]);
3321 let run_array = builder.finish();
3322
3323 let indices = Int32Array::from(vec![Some(0), None, Some(2), Some(3), Some(4)]);
3324 let taken = take(&run_array, &indices, None).unwrap();
3325 let run = taken.as_run::<Int32Type>();
3326 let logical: Vec<Option<i32>> = run.downcast::<Int32Array>().unwrap().into_iter().collect();
3327 assert_eq!(logical, vec![Some(10), None, None, None, Some(99)]);
3328 assert_eq!(run.run_ends().values(), &[1_i32, 4, 5]);
3329 }
3330
3331 #[test]
3332 fn test_take_runs_null_indices_non_nullable_values() {
3333 let run_array = unsafe {
3334 RunArray::<Int32Type>::new_unchecked(
3335 DataType::RunEndEncoded(
3336 Arc::new(Field::new("run_ends", DataType::Int32, false)),
3337 Arc::new(Field::new("values", DataType::Int32, false)),
3338 ),
3339 RunEndBuffer::new(vec![1_i32, 2].into(), 0, 2),
3340 Arc::new(Int32Array::from(vec![10, 20])),
3341 )
3342 };
3343 let indices = Int32Array::from(vec![Some(0), None]);
3344
3345 let error = take(&run_array, &indices, None).unwrap_err();
3346 assert_eq!(
3347 error.to_string(),
3348 "Compute error: Cannot take null indices from a RunEndEncoded array with a non-nullable values field"
3349 );
3350 }
3351
3352 #[test]
3353 fn test_take_runs_sliced() {
3354 let logical_array: Vec<i32> = vec![1, 1, 2, 2, 3, 3, 3, 4, 4, 5, 5, 6, 6];
3355
3356 let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
3357 builder.extend(logical_array.into_iter().map(Some));
3358 let run_array = builder.finish();
3359
3360 let run_array = run_array.slice(4, 6); let take_indices: PrimitiveArray<Int32Type> = vec![0, 5, 5, 1, 4].into_iter().collect();
3363
3364 let result = take_run(&run_array, &take_indices).unwrap();
3365 let result = result.downcast::<Int32Array>().unwrap();
3366
3367 assert_eq!(
3369 result.run_ends().values().len(),
3370 4,
3371 "expected four physical runs"
3372 );
3373 assert_eq!(result.run_ends().values(), &[1_i32, 3, 4, 5]);
3374
3375 let expected = vec![3, 5, 5, 3, 4];
3376 let actual = result.into_iter().flatten().collect::<Vec<_>>();
3377
3378 assert_eq!(expected, actual);
3379 }
3380
3381 #[test]
3382 fn test_take_value_index_from_fixed_list() {
3383 let list = FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>(
3384 vec![
3385 Some(vec![Some(1), Some(2), None]),
3386 Some(vec![Some(4), None, Some(6)]),
3387 None,
3388 Some(vec![None, Some(8), Some(9)]),
3389 ],
3390 3,
3391 );
3392
3393 let indices = UInt32Array::from(vec![2, 1, 0]);
3394 let indexed = take_value_indices_from_fixed_size_list(&list, &indices, 3).unwrap();
3395
3396 assert_eq!(indexed, UInt32Array::from(vec![6, 7, 8, 3, 4, 5, 0, 1, 2]));
3397
3398 let indices = UInt32Array::from(vec![3, 2, 1, 2, 0]);
3399 let indexed = take_value_indices_from_fixed_size_list(&list, &indices, 3).unwrap();
3400
3401 assert_eq!(
3402 indexed,
3403 UInt32Array::from(vec![9, 10, 11, 6, 7, 8, 3, 4, 5, 6, 7, 8, 0, 1, 2])
3404 );
3405 }
3406
3407 #[test]
3408 fn test_take_null_indices() {
3409 let indices = Int32Array::new(
3411 vec![1, 2, 400, 400].into(),
3412 Some(NullBuffer::from(vec![true, true, false, false])),
3413 );
3414 let values = Int32Array::from(vec![1, 23, 4, 5]);
3415 let r = take(&values, &indices, None).unwrap();
3416 let values = r
3417 .as_primitive::<Int32Type>()
3418 .into_iter()
3419 .collect::<Vec<_>>();
3420 assert_eq!(&values, &[Some(23), Some(4), None, None])
3421 }
3422
3423 #[test]
3424 fn test_take_fixed_size_list_null_indices() {
3425 let indices = Int32Array::from_iter([Some(0), None]);
3426 let values = Arc::new(Int32Array::from(vec![0, 1, 2, 3]));
3427 let arr_field = Arc::new(Field::new_list_field(values.data_type().clone(), true));
3428 let values = FixedSizeListArray::try_new(arr_field, 2, values, None).unwrap();
3429
3430 let r = take(&values, &indices, None).unwrap();
3431 let values = r
3432 .as_fixed_size_list()
3433 .values()
3434 .as_primitive::<Int32Type>()
3435 .into_iter()
3436 .collect::<Vec<_>>();
3437 assert_eq!(values, &[Some(0), Some(1), None, None])
3438 }
3439
3440 #[test]
3441 fn test_take_bytes_null_indices() {
3442 let indices = Int32Array::new(
3443 vec![0, 1, 400, 400].into(),
3444 Some(NullBuffer::from_iter(vec![true, true, false, false])),
3445 );
3446 let values = StringArray::from(vec![Some("foo"), None]);
3447 let r = take(&values, &indices, None).unwrap();
3448 let values = r.as_string::<i32>().iter().collect::<Vec<_>>();
3449 assert_eq!(&values, &[Some("foo"), None, None, None])
3450 }
3451
3452 #[test]
3453 fn test_take_union_sparse() {
3454 let structs = create_test_struct(vec![
3455 Some((Some(true), Some(42))),
3456 Some((Some(false), Some(28))),
3457 Some((Some(false), Some(19))),
3458 Some((Some(true), Some(31))),
3459 None,
3460 ]);
3461 let strings = StringArray::from(vec![Some("a"), None, Some("c"), None, Some("d")]);
3462 let type_ids = [1; 5].into_iter().collect::<ScalarBuffer<i8>>();
3463
3464 let union_fields = [
3465 (
3466 0,
3467 Arc::new(Field::new("f1", structs.data_type().clone(), true)),
3468 ),
3469 (
3470 1,
3471 Arc::new(Field::new("f2", strings.data_type().clone(), true)),
3472 ),
3473 ]
3474 .into_iter()
3475 .collect();
3476 let children = vec![Arc::new(structs) as Arc<dyn Array>, Arc::new(strings)];
3477 let array = UnionArray::try_new(union_fields, type_ids, None, children).unwrap();
3478
3479 let indices = vec![0, 3, 1, 0, 2, 4];
3480 let index = UInt32Array::from(indices.clone());
3481 let actual = take(&array, &index, None).unwrap();
3482 let actual = actual.as_any().downcast_ref::<UnionArray>().unwrap();
3483 let strings = actual.child(1);
3484 let strings = strings.as_any().downcast_ref::<StringArray>().unwrap();
3485
3486 let actual = strings.iter().collect::<Vec<_>>();
3487 let expected = vec![Some("a"), None, None, Some("a"), Some("c"), Some("d")];
3488 assert_eq!(expected, actual);
3489 }
3490
3491 #[test]
3492 fn test_take_union_dense() {
3493 let type_ids = vec![0, 1, 1, 0, 0, 1, 0];
3494 let offsets = vec![0, 0, 1, 1, 2, 2, 3];
3495 let ints = vec![10, 20, 30, 40];
3496 let strings = vec![Some("a"), None, Some("c"), Some("d")];
3497
3498 let indices = vec![0, 3, 1, 0, 2, 4];
3499
3500 let taken_type_ids = vec![0, 0, 1, 0, 1, 0];
3501 let taken_offsets = vec![0, 1, 0, 2, 1, 3];
3502 let taken_ints = vec![10, 20, 10, 30];
3503 let taken_strings = vec![Some("a"), None];
3504
3505 let type_ids = <ScalarBuffer<i8>>::from(type_ids);
3506 let offsets = <ScalarBuffer<i32>>::from(offsets);
3507 let ints = UInt32Array::from(ints);
3508 let strings = StringArray::from(strings);
3509
3510 let union_fields = [
3511 (
3512 0,
3513 Arc::new(Field::new("f1", ints.data_type().clone(), true)),
3514 ),
3515 (
3516 1,
3517 Arc::new(Field::new("f2", strings.data_type().clone(), true)),
3518 ),
3519 ]
3520 .into_iter()
3521 .collect();
3522
3523 let array = UnionArray::try_new(
3524 union_fields,
3525 type_ids,
3526 Some(offsets),
3527 vec![Arc::new(ints), Arc::new(strings)],
3528 )
3529 .unwrap();
3530
3531 let index = UInt32Array::from(indices);
3532
3533 let actual = take(&array, &index, None).unwrap();
3534 let actual = actual.as_any().downcast_ref::<UnionArray>().unwrap();
3535
3536 assert_eq!(actual.offsets(), Some(&ScalarBuffer::from(taken_offsets)));
3537 assert_eq!(actual.type_ids(), &ScalarBuffer::from(taken_type_ids));
3538 assert_eq!(
3539 UInt32Array::from(actual.child(0).to_data()),
3540 UInt32Array::from(taken_ints)
3541 );
3542 assert_eq!(
3543 StringArray::from(actual.child(1).to_data()),
3544 StringArray::from(taken_strings)
3545 );
3546 }
3547
3548 fn union_i32_logical(array: &UnionArray) -> Vec<Option<i32>> {
3549 (0..array.len())
3550 .map(|i| {
3551 let child = array.child(array.type_id(i)).as_primitive::<Int32Type>();
3552 let offset = array.value_offset(i);
3553 if child.is_null(offset) {
3554 None
3555 } else {
3556 Some(child.value(offset))
3557 }
3558 })
3559 .collect()
3560 }
3561
3562 #[test]
3563 fn test_take_union_dense_null_indices() {
3564 let fields =
3567 UnionFields::try_new(vec![5], vec![Field::new("i", DataType::Int32, true)]).unwrap();
3568 let dense = UnionArray::try_new(
3569 fields.clone(),
3570 ScalarBuffer::from(vec![5_i8, 5, 5]),
3571 Some(ScalarBuffer::from(vec![0_i32, 1, 2])),
3572 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
3573 )
3574 .unwrap();
3575 let sparse = UnionArray::try_new(
3576 fields,
3577 ScalarBuffer::from(vec![5_i8, 5, 5]),
3578 None,
3579 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
3580 )
3581 .unwrap();
3582
3583 let indices = UInt32Array::new(
3586 ScalarBuffer::from(vec![0_u32, 99, 2]),
3587 Some(NullBuffer::from(vec![true, false, true])),
3588 );
3589 let dense_taken = take(&dense, &indices, None).unwrap();
3590 let sparse_taken = take(&sparse, &indices, None).unwrap();
3591 let dense_logical = union_i32_logical(dense_taken.as_any().downcast_ref().unwrap());
3592 let sparse_logical = union_i32_logical(sparse_taken.as_any().downcast_ref().unwrap());
3593
3594 assert_eq!(dense_logical, vec![Some(1), None, Some(3)]);
3595 assert_eq!(dense_logical, sparse_logical);
3596 }
3597
3598 #[test]
3599 fn test_take_dense_union_null_indices_uses_nullable_field() {
3600 let fields = UnionFields::try_new(
3601 vec![0, 1],
3602 vec![
3603 Field::new("non_nullable", DataType::Int32, false),
3604 Field::new("nullable", DataType::Int32, true),
3605 ],
3606 )
3607 .unwrap();
3608 let union = UnionArray::try_new(
3609 fields,
3610 ScalarBuffer::from(vec![0_i8]),
3611 Some(ScalarBuffer::from(vec![0_i32])),
3612 vec![
3613 Arc::new(Int32Array::from(vec![10])),
3614 Arc::new(Int32Array::from(vec![20])),
3615 ],
3616 )
3617 .unwrap();
3618
3619 let taken = take(&union, &UInt32Array::from(vec![None]), None).unwrap();
3620 let taken = taken.as_union();
3621 assert_eq!(taken.type_id(0), 1);
3622 assert!(taken.child(1).is_null(0));
3623 }
3624
3625 #[test]
3626 fn test_take_sparse_union_null_indices_uses_nullable_field() {
3627 let fields = UnionFields::try_new(
3628 vec![0, 1],
3629 vec![
3630 Field::new("non_nullable", DataType::Int32, false),
3631 Field::new("nullable", DataType::Int32, true),
3632 ],
3633 )
3634 .unwrap();
3635 let union = UnionArray::try_new(
3636 fields,
3637 ScalarBuffer::from(vec![0_i8]),
3638 None,
3639 vec![
3640 Arc::new(Int32Array::from(vec![10])),
3641 Arc::new(Int32Array::from(vec![20])),
3642 ],
3643 )
3644 .unwrap();
3645
3646 let taken = take(&union, &UInt32Array::from(vec![Some(0), None]), None).unwrap();
3647 let taken = taken.as_union();
3648 assert_eq!(taken.type_ids(), &ScalarBuffer::from(vec![0_i8, 1]));
3649 assert_eq!(taken.logical_null_count(), 1);
3650 assert_eq!(
3651 taken
3652 .child(0)
3653 .as_primitive::<Int32Type>()
3654 .iter()
3655 .collect::<Vec<_>>(),
3656 vec![Some(10), Some(10)]
3657 );
3658 assert_eq!(
3659 taken
3660 .child(1)
3661 .as_primitive::<Int32Type>()
3662 .iter()
3663 .collect::<Vec<_>>(),
3664 vec![Some(20), None]
3665 );
3666 }
3667
3668 #[test]
3669 fn test_take_sparse_union_null_indices_no_nullable_fields() {
3670 let fields = UnionFields::try_new(
3671 vec![0, 1],
3672 vec![
3673 Field::new("a", DataType::Int32, false),
3674 Field::new("b", DataType::Int32, false),
3675 ],
3676 )
3677 .unwrap();
3678 let union = UnionArray::try_new(
3679 fields,
3680 ScalarBuffer::from(vec![0_i8]),
3681 None,
3682 vec![
3683 Arc::new(Int32Array::from(vec![10])),
3684 Arc::new(Int32Array::from(vec![20])),
3685 ],
3686 )
3687 .unwrap();
3688
3689 let error = take(&union, &UInt32Array::from(vec![None]), None).unwrap_err();
3690 assert_eq!(
3691 error.to_string(),
3692 "Compute error: Cannot take null indices from a union with no field that can represent nulls"
3693 );
3694 }
3695
3696 #[test]
3697 fn test_take_empty_sparse_union_null_indices_mixed_nullability() {
3698 let fields = UnionFields::try_new(
3699 vec![0, 1],
3700 vec![
3701 Field::new("non_nullable", DataType::Int32, false),
3702 Field::new("nullable", DataType::Int32, true),
3703 ],
3704 )
3705 .unwrap();
3706 let union = UnionArray::try_new(
3707 fields,
3708 ScalarBuffer::<i8>::from(vec![]),
3709 None,
3710 vec![
3711 Arc::new(Int32Array::from(Vec::<i32>::new())),
3712 Arc::new(Int32Array::from(Vec::<i32>::new())),
3713 ],
3714 )
3715 .unwrap();
3716
3717 let taken = take(&union, &UInt32Array::from(vec![None]), None).unwrap();
3718 let taken = taken.as_union();
3719 assert_eq!(taken.len(), 1);
3720 assert_eq!(taken.type_id(0), 1);
3721 assert_eq!(taken.logical_null_count(), 1);
3722 assert!(!taken.child(0).is_null(0));
3723 assert!(taken.child(1).is_null(0));
3724 }
3725
3726 #[test]
3727 fn test_take_empty_union_without_null_indices() {
3728 let fields = UnionFields::try_new(vec![], Vec::<Field>::new()).unwrap();
3729 let indices = UInt32Array::from(Vec::<u32>::new());
3730
3731 let sparse = UnionArray::try_new(
3732 fields.clone(),
3733 ScalarBuffer::<i8>::from(vec![]),
3734 None,
3735 vec![],
3736 )
3737 .unwrap();
3738 let dense = UnionArray::try_new(
3739 fields,
3740 ScalarBuffer::<i8>::from(vec![]),
3741 Some(ScalarBuffer::<i32>::from(vec![])),
3742 vec![],
3743 )
3744 .unwrap();
3745
3746 for values in [&sparse, &dense] {
3747 let taken = take(values, &indices, None).unwrap();
3748 assert_eq!(taken.len(), 0);
3749 assert_eq!(taken.data_type(), values.data_type());
3750 }
3751 }
3752
3753 #[test]
3754 fn test_take_empty_union_with_null_indices() {
3755 let fields = UnionFields::try_new(vec![], Vec::<Field>::new()).unwrap();
3756 let indices = UInt32Array::from(vec![None]);
3757
3758 let sparse = UnionArray::try_new(
3759 fields.clone(),
3760 ScalarBuffer::<i8>::from(vec![]),
3761 None,
3762 vec![],
3763 )
3764 .unwrap();
3765 let dense = UnionArray::try_new(
3766 fields,
3767 ScalarBuffer::<i8>::from(vec![]),
3768 Some(ScalarBuffer::<i32>::from(vec![])),
3769 vec![],
3770 )
3771 .unwrap();
3772
3773 for values in [&sparse, &dense] {
3774 let error = take(values, &indices, None).unwrap_err();
3775 assert_eq!(
3776 error.to_string(),
3777 "Compute error: Cannot take null indices from a union with no field that can represent nulls"
3778 );
3779 }
3780 }
3781
3782 fn inner_union_no_nullable_children(mode: UnionMode) -> (UnionFields, UnionArray) {
3783 let fields = UnionFields::try_new(
3784 vec![0, 1],
3785 vec![
3786 Field::new("a", DataType::Int32, false),
3787 Field::new("b", DataType::Int32, false),
3788 ],
3789 )
3790 .unwrap();
3791 let children: Vec<ArrayRef> = vec![
3792 Arc::new(Int32Array::from(vec![10])),
3793 Arc::new(Int32Array::from(vec![20])),
3794 ];
3795 let array = match mode {
3796 UnionMode::Sparse => UnionArray::try_new(
3797 fields.clone(),
3798 ScalarBuffer::from(vec![0_i8]),
3799 None,
3800 children,
3801 )
3802 .unwrap(),
3803 UnionMode::Dense => UnionArray::try_new(
3804 fields.clone(),
3805 ScalarBuffer::from(vec![0_i8]),
3806 Some(ScalarBuffer::from(vec![0_i32])),
3807 children,
3808 )
3809 .unwrap(),
3810 };
3811 (fields, array)
3812 }
3813
3814 #[test]
3815 fn test_take_union_null_indices_skips_nested_union_without_nullable_children() {
3816 for mode in [UnionMode::Sparse, UnionMode::Dense] {
3820 let (inner_fields, inner) = inner_union_no_nullable_children(mode);
3821 let outer_fields = UnionFields::try_new(
3822 vec![0, 1],
3823 vec![
3824 Field::new("inner", DataType::Union(inner_fields, mode), true),
3825 Field::new("i", DataType::Int32, true),
3826 ],
3827 )
3828 .unwrap();
3829 let children: Vec<ArrayRef> =
3830 vec![Arc::new(inner), Arc::new(Int32Array::from(vec![30]))];
3831 let outer = match mode {
3832 UnionMode::Sparse => UnionArray::try_new(
3833 outer_fields,
3834 ScalarBuffer::from(vec![1_i8]),
3835 None,
3836 children,
3837 )
3838 .unwrap(),
3839 UnionMode::Dense => UnionArray::try_new(
3840 outer_fields,
3841 ScalarBuffer::from(vec![1_i8]),
3842 Some(ScalarBuffer::from(vec![0_i32])),
3843 children,
3844 )
3845 .unwrap(),
3846 };
3847
3848 let taken = take(&outer, &UInt32Array::from(vec![None]), None).unwrap();
3849 let taken = taken.as_union();
3850 assert_eq!(taken.type_id(0), 1, "{mode:?}");
3851 assert!(taken.child(1).is_null(0), "{mode:?}");
3852 assert_eq!(taken.logical_null_count(), 1, "{mode:?}");
3853 }
3854 }
3855
3856 #[test]
3857 fn test_take_union_null_indices_nested_union_cannot_represent_null() {
3858 for mode in [UnionMode::Sparse, UnionMode::Dense] {
3861 let (inner_fields, inner) = inner_union_no_nullable_children(mode);
3862 let outer_fields = UnionFields::try_new(
3863 vec![0],
3864 vec![Field::new(
3865 "inner",
3866 DataType::Union(inner_fields, mode),
3867 true,
3868 )],
3869 )
3870 .unwrap();
3871 let children: Vec<ArrayRef> = vec![Arc::new(inner)];
3872 let outer = match mode {
3873 UnionMode::Sparse => UnionArray::try_new(
3874 outer_fields,
3875 ScalarBuffer::from(vec![0_i8]),
3876 None,
3877 children,
3878 )
3879 .unwrap(),
3880 UnionMode::Dense => UnionArray::try_new(
3881 outer_fields,
3882 ScalarBuffer::from(vec![0_i8]),
3883 Some(ScalarBuffer::from(vec![0_i32])),
3884 children,
3885 )
3886 .unwrap(),
3887 };
3888
3889 let error = take(&outer, &UInt32Array::from(vec![None]), None).unwrap_err();
3890 assert_eq!(
3891 error.to_string(),
3892 "Compute error: Cannot take null indices from a union with no field that can represent nulls"
3893 );
3894 }
3895 }
3896
3897 #[test]
3898 fn test_take_union_null_indices_skips_ree_with_non_nullable_values() {
3899 let run_array = unsafe {
3902 RunArray::<Int32Type>::new_unchecked(
3903 DataType::RunEndEncoded(
3904 Arc::new(Field::new("run_ends", DataType::Int32, false)),
3905 Arc::new(Field::new("values", DataType::Int32, false)),
3906 ),
3907 RunEndBuffer::new(vec![1_i32].into(), 0, 1),
3908 Arc::new(Int32Array::from(vec![10])),
3909 )
3910 };
3911 let fields = UnionFields::try_new(
3912 vec![0, 1],
3913 vec![
3914 Field::new("ree", run_array.data_type().clone(), true),
3915 Field::new("i", DataType::Int32, true),
3916 ],
3917 )
3918 .unwrap();
3919 let union = UnionArray::try_new(
3920 fields,
3921 ScalarBuffer::from(vec![1_i8]),
3922 None,
3923 vec![
3924 Arc::new(run_array) as ArrayRef,
3925 Arc::new(Int32Array::from(vec![30])),
3926 ],
3927 )
3928 .unwrap();
3929
3930 let taken = take(&union, &UInt32Array::from(vec![None]), None).unwrap();
3931 let taken = taken.as_union();
3932 assert_eq!(taken.type_id(0), 1);
3933 assert!(taken.child(1).is_null(0));
3934 assert_eq!(taken.logical_null_count(), 1);
3935 }
3936
3937 #[test]
3938 fn test_take_union_dense_using_builder() {
3939 let mut builder = UnionBuilder::new_dense();
3940
3941 builder.append::<Int32Type>("a", 1).unwrap();
3942 builder.append::<Float64Type>("b", 3.0).unwrap();
3943 builder.append::<Int32Type>("a", 4).unwrap();
3944 builder.append::<Int32Type>("a", 5).unwrap();
3945 builder.append::<Float64Type>("b", 2.0).unwrap();
3946
3947 let union = builder.build().unwrap();
3948
3949 let indices = UInt32Array::from(vec![2, 0, 1, 2]);
3950
3951 let mut builder = UnionBuilder::new_dense();
3952
3953 builder.append::<Int32Type>("a", 4).unwrap();
3954 builder.append::<Int32Type>("a", 1).unwrap();
3955 builder.append::<Float64Type>("b", 3.0).unwrap();
3956 builder.append::<Int32Type>("a", 4).unwrap();
3957
3958 let taken = builder.build().unwrap();
3959
3960 assert_eq!(
3961 taken.to_data(),
3962 take(&union, &indices, None).unwrap().to_data()
3963 );
3964 }
3965
3966 #[test]
3967 fn test_take_union_dense_all_match_issue_6206() {
3968 let fields = UnionFields::from_fields(vec![Field::new("a", DataType::Int64, false)]);
3969 let ints = Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5]));
3970
3971 let array = UnionArray::try_new(
3972 fields,
3973 ScalarBuffer::from(vec![0_i8, 0, 0, 0, 0]),
3974 Some(ScalarBuffer::from_iter(0_i32..5)),
3975 vec![ints],
3976 )
3977 .unwrap();
3978
3979 let indices = Int64Array::from(vec![0, 2, 4]);
3980 let array = take(&array, &indices, None).unwrap();
3981 assert_eq!(array.len(), 3);
3982 }
3983
3984 fn offset_overflow_fixture() -> (StringArray, usize) {
3989 let value_len = 1_000_000;
3990 let values = StringArray::from(vec![Some("a".repeat(value_len))]);
3991 let n = i32::MAX as usize / value_len + 1;
3992 (values, n)
3993 }
3994
3995 #[test]
3996 fn test_take_bytes_offset_overflow() {
3997 let (values, n) = offset_overflow_fixture();
3998 let indices = Int32Array::from(vec![0; n]);
3999 assert!(matches!(
4000 take(&values, &indices, None),
4001 Err(ArrowError::OffsetOverflowError(_))
4002 ));
4003 }
4004
4005 #[test]
4008 fn test_take_bytes_offset_overflow_nullable() {
4009 let (values, n) = offset_overflow_fixture();
4010 let validity =
4013 NullBuffer::from_iter(std::iter::once(false).chain(std::iter::repeat_n(true, n)));
4014 let indices = Int32Array::new(vec![0i32; n + 1].into(), Some(validity));
4015
4016 assert!(matches!(
4017 take(&values, &indices, None),
4018 Err(ArrowError::OffsetOverflowError(_))
4019 ));
4020 }
4021
4022 #[test]
4023 fn test_take_run_empty_indices() {
4024 let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
4025 builder.extend([Some(1), Some(1), Some(2), Some(2)]);
4026 let run_array = builder.finish();
4027
4028 let logical_indices: PrimitiveArray<Int32Type> = PrimitiveArray::from(Vec::<i32>::new());
4029
4030 let result = take_impl::<_, true>(&run_array, &logical_indices)
4031 .expect("take_run with empty indices");
4032
4033 assert_eq!(result.len(), 0);
4035 assert_eq!(result.null_count(), 0);
4036
4037 let run_result = result
4040 .as_any()
4041 .downcast_ref::<RunArray<Int32Type>>()
4042 .expect("result should be a RunArray");
4043 assert_eq!(run_result.run_ends().len(), 0);
4044 assert_eq!(run_result.values().len(), 0);
4045 }
4046
4047 #[test]
4048 fn test_take_run_end_encoded_merges_identical_runs() {
4049 let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
4053 builder.extend([1, 1, 0, 0, 1, 1].into_iter().map(Some));
4054 let ree = builder.finish();
4055
4056 let indexes = Int32Array::from_iter_values(vec![0, 1, 4, 5]);
4057 let result = take(&ree, &indexes, None).unwrap();
4058 let result = result
4059 .as_run::<Int32Type>()
4060 .downcast::<Int32Array>()
4061 .unwrap();
4062
4063 assert_eq!(
4065 result.run_ends().values().len(),
4066 1,
4067 "expected a single physical run"
4068 );
4069 assert_eq!(result.run_ends().values(), &[4_i32]);
4070
4071 let actual = result.into_iter().flatten().collect::<Vec<_>>();
4072 assert_eq!(actual, vec![1, 1, 1, 1]);
4073 }
4074
4075 #[test]
4076 fn test_take_run_end_encoded_merges_identical_string_runs() {
4077 let mut builder = StringRunBuilder::<Int32Type>::new();
4078 builder.extend(
4079 ["bob", "bob", "alice", "alice", "bob", "bob"]
4080 .into_iter()
4081 .map(Some),
4082 );
4083 let ree = builder.finish();
4084
4085 let indexes = Int32Array::from_iter_values(vec![0, 1, 4, 5]);
4086 let result = take(&ree, &indexes, None).unwrap();
4087 let result = result
4088 .as_run::<Int32Type>()
4089 .downcast::<StringArray>()
4090 .unwrap();
4091
4092 assert_eq!(
4094 result.run_ends().values().len(),
4095 1,
4096 "expected a single physical run"
4097 );
4098 assert_eq!(result.run_ends().values(), &[4_i32]);
4099
4100 let actual = result.into_iter().flatten().collect::<Vec<_>>();
4101 assert_eq!(actual, vec!["bob", "bob", "bob", "bob"]);
4102 }
4103
4104 #[test]
4105 fn test_take_run_end_encoded_mixed_runs() {
4106 let mut builder = StringRunBuilder::<Int32Type>::new();
4109 builder.extend(
4110 ["bob", "bob", "alice", "alice", "bob", "bob", "eve", "eve"]
4111 .into_iter()
4112 .map(Some),
4113 );
4114 let ree = builder.finish();
4115
4116 let indexes = Int32Array::from_iter_values(vec![0, 0, 1, 4, 5, 2, 3, 2, 6, 7, 6]);
4118 let result = take(&ree, &indexes, None).unwrap();
4119 let result = result
4120 .as_run::<Int32Type>()
4121 .downcast::<StringArray>()
4122 .unwrap();
4123
4124 println!("run_ends_raw: {:?}", result.run_ends());
4127 println!("run_ends: {:?}", result.run_ends().values());
4128 println!("values : {:?}", result.values());
4129 assert_eq!(
4130 result.run_ends().values().len(),
4131 3,
4132 "expected three physical runs"
4133 );
4134 assert_eq!(result.run_ends().values(), &[5_i32, 8, 11]);
4135
4136 let actual = result.into_iter().flatten().collect::<Vec<_>>();
4137 assert_eq!(
4138 actual,
4139 vec![
4140 "bob", "bob", "bob", "bob", "bob", "alice", "alice", "alice", "eve", "eve", "eve"
4141 ]
4142 );
4143 }
4144
4145 #[test]
4147 fn test_take_fixed_size_list_parent_nulls() {
4148 let list = FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>(
4149 vec![
4150 Some(vec![Some(1), Some(2)]),
4151 None,
4152 Some(vec![Some(5), Some(6)]),
4153 ],
4154 2,
4155 );
4156
4157 let indices = UInt32Array::from(vec![2, 1, 0]);
4158 let result = take(&list, &indices, None).unwrap();
4159 let result = result.as_fixed_size_list();
4160
4161 assert_eq!(result.len(), 3);
4162 assert!(result.is_valid(0));
4163 assert!(result.is_null(1));
4164 assert!(result.is_valid(2));
4165
4166 let child = result.values().as_primitive::<Int32Type>();
4167 assert_eq!(child.value(0), 5);
4168 assert_eq!(child.value(1), 6);
4169 assert!(child.is_null(2));
4170 assert!(child.is_null(3));
4171 assert_eq!(child.value(4), 1);
4172 assert_eq!(child.value(5), 2);
4173 }
4174
4175 #[test]
4176 fn test_take_zero_sized_fixed_size_list() {
4177 let input = FixedSizeListArray::try_new_with_length(
4178 Field::new_list_field(DataType::Int32, true).into(),
4179 0,
4180 Arc::new(Int32Array::new_null(0)),
4181 None,
4182 3,
4183 )
4184 .unwrap();
4185
4186 let indices = UInt32Array::from(vec![2, 0]);
4187 let result = take(&input, &indices, None).unwrap();
4188
4189 assert_eq!(result.len(), 2);
4190 }
4191}