Skip to main content

arrow_select/
take.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Defines take kernel for [Array]
19
20use 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
38/// Take elements by index from [Array], creating a new [Array] from those indexes.
39///
40/// ```text
41/// ┌─────────────────┐      ┌─────────┐                              ┌─────────────────┐
42/// │        A        │      │    0    │                              │        A        │
43/// ├─────────────────┤      ├─────────┤                              ├─────────────────┤
44/// │        D        │      │    2    │                              │        B        │
45/// ├─────────────────┤      ├─────────┤   take(values, indices)      ├─────────────────┤
46/// │        B        │      │    3    │ ─────────────────────────▶   │        C        │
47/// ├─────────────────┤      ├─────────┤                              ├─────────────────┤
48/// │        C        │      │    1    │                              │        D        │
49/// ├─────────────────┤      └─────────┘                              └─────────────────┘
50/// │        E        │
51/// └─────────────────┘
52///    values array          indices array                              result
53/// ```
54///
55/// For selecting values by index from multiple arrays see [`crate::interleave`]
56///
57/// Note that this kernel, similar to other kernels in this crate,
58/// will avoid allocating where not necessary. Consequently
59/// the returned array may share buffers with the inputs
60///
61/// # Errors
62/// This function errors whenever:
63/// * An index cannot be casted to `usize` (typically 32 bit architectures)
64/// * An index is out of bounds and `options` is set to check bounds.
65///
66/// # Panics
67///
68/// When `options` is not set to check bounds, taking indexes after `len` will panic.
69///
70/// # See also
71/// * [`BatchCoalescer`]: to filter multiple [`RecordBatch`] and coalesce
72///   the results into a single array.
73///
74/// [`BatchCoalescer`]: crate::coalesce::BatchCoalescer
75///
76/// # Examples
77/// ```
78/// # use arrow_array::{StringArray, UInt32Array, cast::AsArray};
79/// # use arrow_select::take::take;
80/// let values = StringArray::from(vec!["zero", "one", "two"]);
81///
82/// // Take items at index 2, and 1:
83/// let indices = UInt32Array::from(vec![2, 1]);
84/// let taken = take(&values, &indices, None).unwrap();
85/// let taken = taken.as_string::<i32>();
86///
87/// assert_eq!(*taken, StringArray::from(vec!["two", "one"]));
88/// ```
89pub 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
107/// For each [ArrayRef] in the [`Vec<ArrayRef>`], take elements by index and create a new
108/// [`Vec<ArrayRef>`] from those indices.
109///
110/// ```text
111/// ┌────────┬────────┐
112/// │        │        │           ┌────────┐                                ┌────────┬────────┐
113/// │   A    │   1    │           │        │                                │        │        │
114/// ├────────┼────────┤           │   0    │                                │   A    │   1    │
115/// │        │        │           ├────────┤                                ├────────┼────────┤
116/// │   D    │   4    │           │        │                                │        │        │
117/// ├────────┼────────┤           │   2    │  take_arrays(values,indices)   │   B    │   2    │
118/// │        │        │           ├────────┤                                ├────────┼────────┤
119/// │   B    │   2    │           │        │  ───────────────────────────►  │        │        │
120/// ├────────┼────────┤           │   3    │                                │   C    │   3    │
121/// │        │        │           ├────────┤                                ├────────┼────────┤
122/// │   C    │   3    │           │        │                                │        │        │
123/// ├────────┼────────┤           │   1    │                                │   D    │   4    │
124/// │        │        │           └────────┘                                └────────┼────────┘
125/// │   E    │   5    │
126/// └────────┴────────┘
127///    values arrays             indices array                                      result
128/// ```
129///
130/// # Errors
131/// This function errors whenever:
132/// * An index cannot be casted to `usize` (typically 32 bit architectures)
133/// * An index is out of bounds and `options` is set to check bounds.
134///
135/// # Panics
136///
137/// When `options` is not set to check bounds, taking indexes after `len` will panic.
138///
139/// # Examples
140/// ```
141/// # use std::sync::Arc;
142/// # use arrow_array::{StringArray, UInt32Array, cast::AsArray};
143/// # use arrow_select::take::{take, take_arrays};
144/// let string_values = Arc::new(StringArray::from(vec!["zero", "one", "two"]));
145/// let values = Arc::new(UInt32Array::from(vec![0, 1, 2]));
146///
147/// // Take items at index 2, and 1:
148/// let indices = UInt32Array::from(vec![2, 1]);
149/// let taken_arrays = take_arrays(&[string_values, values], &indices, None).unwrap();
150/// let taken_string = taken_arrays[0].as_string::<i32>();
151/// assert_eq!(*taken_string, StringArray::from(vec!["two", "one"]));
152/// let taken_values = taken_arrays[1].as_primitive();
153/// assert_eq!(*taken_values, UInt32Array::from(vec![2, 1]));
154/// ```
155pub 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
166/// For each [ArrayRef] in the [`Vec<ArrayRef>`], take elements by index and create a new
167/// [`Vec<ArrayRef>`] from those indices, without bounds checking.
168///
169/// # Safety
170///
171/// Caller must ensure all values in `indices` are valid indices into each array
172/// (i.e. `< array.len()`). Out-of-bounds indices will cause undefined behavior.
173///
174/// # Errors
175///
176/// Returns an error if an index value cannot be cast to `usize`.
177pub 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
193/// Verifies that the non-null values of `indices` are all `< len`
194fn 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            // the biggest representable value for T::Native is lower than len, e.g: u8::MAX < 512, no need to check bounds
204            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            // `new_empty_array` cannot construct a union with no fields, but an existing empty
248            // union can be sliced without materializing a child array.
249            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            // Create the null bit buffer.
315            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            // Take applied to a null array produces a null array.
359            if values.len() >= indices.len() {
360                // If the existing null array is as big as the indices, we can use a slice of it
361                // to avoid allocating a new null array.
362                Ok(values.slice(0, indices.len()))
363            } else {
364                // If the existing null array isn't big enough, create a new one.
365                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            // Keep index nulls so `take` of each child writes a null instead of
391            // reading child offset 0 (the default `take_native` fills in).
392            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
432/// Takes union type ids, substituting a child that can represent a null for
433/// null take indices.
434///
435/// Union arrays do not have a top-level null bitmap. A null is represented by
436/// selecting a child that can store a null value. In particular, a null index
437/// cannot fall back to type id `0`, as unions are not required to have such a
438/// child.
439///
440/// Returns the taken type ids and, when `indices` contains nulls, the type id
441/// used to represent those nulls.
442fn 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
467/// Type id of a union child that can represent a newly introduced null.
468fn 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
480/// Whether `field` can physically store a newly introduced null.
481///
482/// Union and RunEndEncoded have no top-level validity bitmap, so a null must
483/// be represented by a descendant that is marked nullable. A field marked
484/// nullable is not sufficient if its nested type cannot store a null.
485fn 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
498/// Takes a sparse union child for `indices`.
499///
500/// Null take indices select one child that can represent a null
501/// ([`take_union_type_ids`]). Values in the other children at those positions
502/// are unspecified, so they are taken with dummy indices instead of introducing
503/// nulls that would contradict field metadata.
504fn 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        // Dummy index 0 is OOB on an empty child. Sparse children have the same
515        // length as the union, so a non-null index is also OOB and already
516        // panics in `take_native` via [`take_union_type_ids`].
517        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
529/// Replaces null take indices with `0` and drops the null bitmap.
530fn 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/// Options that define how `take` should behave
539#[derive(Clone, Debug, Default)]
540pub struct TakeOptions {
541    /// Perform bounds check before taking indices from values.
542    /// If enabled, an `ArrowError` is returned if the indices are out of bounds.
543    /// If not enabled, and indices exceed bounds, the kernel will panic.
544    pub check_bounds: bool,
545}
546
547/// `take` implementation for all primitive arrays
548///
549/// This checks if an `indices` slot is populated, and gets the value from `values`
550///  as the populated index.
551/// If the `indices` slot is null, a null value is returned.
552/// For example, given:
553///     values:  [1, 2, 3, null, 5]
554///     indices: [0, null, 4, 3]
555/// The result is: [1 (slot 0), null (null slot), 5 (slot 4), null (slot 3)]
556fn 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                // SAFETY: idx<indices.len()
596                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                        // SAFETY: caller guarantees all indices are in-bounds.
615                        unsafe { *values.get_unchecked(index.as_usize()) }
616                    })
617                    .collect()
618            }
619        }
620    }
621}
622
623/// Read the bit at `src_bit_idx` from `src` and, if it is set, write a `1` to `dst_bit_idx`
624/// in `dst`. Leaves `dst_bit_idx` unchanged (zero) when the source bit is unset.
625///
626/// ```text
627/// src = 0b00100000  (bit 5 is set)
628/// copy_bit_if_set(src, 5, dst, 2)  →  dst bit 2 becomes 1
629/// ```
630///
631/// # Safety
632/// - `src` must be valid for reads up to byte `src_bit_idx / 8`.
633/// - `dst` must be valid for writes up to byte `dst_bit_idx / 8`.
634#[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/// Read the bit at `bit_idx` from `src` and return it shifted to `out_pos`, ready to be
644/// OR'd into an output byte accumulator.
645///
646/// ```text
647/// src = 0b10100000  (bit 5 is set)
648/// pack_bit(src, 5, 2)  →  0b00000100   (bit from position 5, placed at position 2)
649/// ```
650///
651/// # Safety
652/// `src` must be valid for reads up to byte `bit_idx / 8`.
653#[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 containing bit `bit_idx`
656    ((byte >> (bit_idx & 7)) & 1) << out_pos // extract the bit, shift to output position
657}
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                // SAFETY: valid_idx < indices.len(), guaranteed by valid_indices().
675                let index_val = unsafe { indices.value_unchecked(valid_idx) }.as_usize();
676                if VALIDATE_INDICES {
677                    if values.value(index_val) {
678                        // SAFETY: valid_idx < indices.len() = len, output buffer holds len bits.
679                        unsafe { bit_util::set_bit_raw(out_ptr, valid_idx) };
680                    }
681                } else {
682                    // SAFETY: caller guarantees index_val < values.len().
683                    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            // Build the output byte-by-byte with an 8-element inner loop so the
690            // compiler can fully unroll it and issue the 8 source loads in parallel.
691            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                    // SAFETY: base + bit < full_bytes * 8 <= len, so base + bit is a valid
700                    // position in the indices array.
701                    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                        // SAFETY: caller guarantees index_val < values.len().
706                        byte |= unsafe { pack_bit(src_ptr, index_val + src_offset, bit) };
707                    }
708                }
709                *out_byte = byte;
710            }
711            // Handle remaining bits when len is not a multiple of 8.
712            if full_bytes < out_bytes {
713                let base = full_bytes * 8;
714                let mut byte = 0u8;
715                for bit in 0..(len - base) {
716                    // SAFETY: base + bit < len (remainder loop bound), so base + bit is a
717                    // valid position in the indices array.
718                    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                        // SAFETY: caller guarantees index_val < values.len().
723                        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/// Gather value bits and validity bits from two boolean buffers in a single pass.
734/// Used when the values array itself has nulls, avoiding two separate `take_bits` calls.
735#[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                // SAFETY: out_pos < indices.len(), guaranteed by valid_indices().
757                let src_idx = unsafe { indices.value_unchecked(out_pos) }.as_usize();
758                if VALIDATE_INDICES {
759                    if values.value(src_idx) {
760                        // SAFETY: out_pos < indices.len() = len, output buffer holds len bits.
761                        unsafe { bit_util::set_bit_raw(value_out_ptr, out_pos) };
762                    }
763                    if validity.value(src_idx) {
764                        // SAFETY: out_pos < indices.len() = len, output buffer holds len bits.
765                        unsafe { bit_util::set_bit_raw(validity_out_ptr, out_pos) };
766                    }
767                } else {
768                    // SAFETY: caller guarantees src_idx < values.len().
769                    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                    // SAFETY: bit_base + bit_pos < full_bytes * 8 <= len.
802                    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                        // SAFETY: caller guarantees src_idx < values.len().
808                        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            // Handle remaining bits when len is not a multiple of 8.
820            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                    // SAFETY: bit_base + bit_pos < len (remainder loop bound).
826                    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                        // SAFETY: caller guarantees src_idx < values.len().
832                        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
851/// `take` implementation for boolean arrays
852fn 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
871/// `take` implementation for string arrays
872fn 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    // Branch on output nulls — `None` means every output slot is valid.
885    match nulls.as_ref().filter(|n| n.null_count() > 0) {
886        // Fast path: no nulls in output, every index is valid.
887        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                // SAFETY: in-bounds proven by the first loop's bounds-checked offset access.
907                // dst asserted above to include the required capacity.
908                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            // SAFETY: wrote exactly `capacity` bytes above; reserved on line above.
920            unsafe {
921                values.set_len(capacity);
922            }
923        }
924        // Nullable path: only process valid (non-null) output positions.
925        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            // Pre-fill offsets; we overwrite valid positions below.
930            offsets.resize(indices.len() + 1, T::Offset::default());
931
932            // Pass 1: find all valid ranges that need to be copied.
933            for i in output_nulls.valid_indices() {
934                let current_offset = T::Offset::from_usize(capacity)
935                    .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?;
936                // Fill offsets for skipped null slots so they get zero-length ranges.
937                if last_filled < i {
938                    offsets[last_filled + 1..=i].fill(current_offset);
939                }
940
941                // SAFETY: `i` comes from a validity bitmap over `indices`, so it is in-bounds.
942                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            // Fill trailing null offsets after the last valid position.
954            let final_offset = T::Offset::from_usize(capacity)
955                .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?;
956            offsets[last_filled + 1..].fill(final_offset);
957            // Pass 2: copy byte data for all collected ranges.
958            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                // SAFETY: caller guarantees each (start, end) is in-bounds of `src`.
975                // `dst` asserted above to include the required capacity.
976                // The regions don't overlap (src is input, dst is a fresh allocation).
977                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            // SAFETY: caller guarantees `capacity` == total bytes across all ranges,
987            // so the loop above wrote exactly `capacity` bytes.
988            unsafe { values.set_len(capacity) };
989        }
990    }
991
992    // SAFETY: offsets are monotonically increasing and in-bounds of `values`,
993    // and `nulls` (if present) has length == `indices.len()`.
994    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
1002/// `take` implementation for byte view arrays
1003fn 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    // Safety:  array.views was valid, and take_native copies only valid values, and verifies bounds
1011    Ok(unsafe { GenericByteViewArray::new_unchecked(new_views, buffers, new_nulls) })
1012}
1013
1014/// `take` implementation for list arrays
1015///
1016/// Copies the selected list entries' child slices into a new child array
1017/// via `MutableArrayData`, then reconstructs a list array with new offsets
1018fn 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    // The number of child values that the rows use, excluding values outside a slice
1031    let visible_child_len = (values.offsets().last() - values.offsets().first()).as_usize();
1032    // An estimate of the output's child length, from the average length of the rows
1033    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                    // Fill offsets for null values between the two valid indices.
1070                    if prev < vidx {
1071                        dst_offsets.extend(std::iter::repeat_n(child_len, vidx - prev));
1072                    }
1073                    // SAFETY: vidx < indices.len(), guaranteed by valid_indices().
1074                    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                        // SAFETY: caller guarantees row < values.len(), so row+1 is also in bounds.
1082                        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        // Safety: data_type, len, and buffer are all derived from the already-validated
1110        // source child_data, so re-validation is unnecessary.
1111        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        // SAFETY: `dst_offsets` is constructed to be monotonically increasing above.
1118        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                // SAFETY: i < indices.len(), guaranteed by valid_indices().
1152                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                    // SAFETY: caller guarantees row < values.len(), so row+1 is also in bounds.
1157                    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            // Filling offsets for null values at the end
1172            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    // SAFETY: `dst_offsets` is constructed to be monotonically increasing above
1181    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    // SAFETY: the taken offsets/sizes are a permutation of the (valid) input
1205    // offsets/sizes, so they remain within the bounds of the child array.
1206    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
1217/// `take` implementation for `FixedSizeListArray`
1218///
1219/// Calculates the index and indexed offset for the inner array,
1220/// applying `take` on the inner array, then reconstructing a list array
1221/// with the indexed offsets
1222fn 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    // Fast path: primitive child with no nulls  copy row-sized byte blocks directly,
1232    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    // Expand the per-row null bitmap to per-element: each null row produces `list_size` null elements.
1294    let child_null_buf = taken_row_nulls.map(|n| n.expand(list_size).buffer().clone());
1295
1296    // SAFETY:
1297    // - The data buffer has `indices.len() * row_bytes` bytes = `out_len` elements of a
1298    //   primitive type, matching `.len(out_len)`.
1299    // - The null buffer (when present) is `ceil(out_len, 8)` bytes via `expand(list_size).buffer()`.
1300    // - Primitive types have exactly one value buffer.
1301    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
1311/// The take kernel implementation for `FixedSizeBinaryArray`.
1312///
1313/// The computation is done in two steps:
1314/// - Compute the values buffer
1315/// - Compute the null buffer
1316fn 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    /// Implementation of the take kernel for fixed size binary arrays.
1340    #[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            // The indices nullability cannot be ignored here because the values buffer may contain
1359            // nulls which should not cause a panic.
1360            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
1378/// Implements the take kernel semantics over a flat [`Buffer`], interpreting it as a slice of
1379/// `&[[u8; N]]`, where `N` is a compile-time constant. The usage of a flat [`Buffer`] allows using
1380/// this kernel without an available [`ArrowPrimitiveType`] (e.g., for `[u8; 5]`).
1381///
1382/// # Using This Function in the Primitive Take Kernel
1383///
1384/// This function is basically the same as [`take_native`] but just on a flat [`Buffer`] instead of
1385/// the primitive [`ScalarBuffer`]. Ideally, the [`take_primitive`] kernel should just use this
1386/// more general function. However, the "idiomatic code" requires the
1387/// [feature(generic_const_exprs)](https://github.com/rust-lang/rust/issues/76560) for calling
1388/// `take_fixed_size<I, { size_of::<T::Native> () } >(...)`. Once this feature has been stabilized,
1389/// we can use this function also in the primitive kernels.
1390fn 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        // SAFETY: interpret an already valid slice as a slice of N-byte chunks. N divides buffer
1405        // length without remainder.
1406        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                // SAFETY: idx<indices.len()
1417                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); // Prevent de-allocation
1431    let ptr = vec.as_mut_ptr();
1432    let len = vec.len();
1433    let cap = vec.capacity();
1434    let result_buffer = unsafe {
1435        // SAFETY: flattening an already valid Vec.
1436        Vec::from_raw_parts(ptr.cast::<u8>(), len * N, cap * N)
1437    };
1438
1439    Buffer::from_vec(result_buffer)
1440}
1441
1442/// `take` implementation for dictionary arrays
1443///
1444/// applies `take` to the keys of the dictionary array and returns a new dictionary array
1445/// with the same dictionary values and reordered keys
1446fn 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
1454/// `take` implementation for run arrays
1455///
1456/// Finds physical indices for the given logical indices and builds output run array
1457/// by taking values in the input run_array.values at the physical indices.
1458/// The output run array will be run encoded on the physical indices and not on output values.
1459/// For e.g. an input `RunArray{ run_ends = [2,4,6,8], values=[1,2,1,2] }` and `logical_indices=[2,3,6,7]`
1460/// would be converted to `physical_indices=[1,1,3,3]` which will be used to build
1461/// output `RunArray{ run_ends=[2,4], values=[2,2] }`.
1462///
1463/// A null logical index becomes a null run. Consecutive nulls are merged.
1464fn 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    // Run encode the physical indices into new_run_ends
1478    // Keep track of the physical indices to take in take_value_indices
1479    // `unwrap` is used in this function because the unwrapped values are bounded by the corresponding `::Native`.
1480    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            // Safe unwrap since physical indices came from a valid run array.
1495            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    // Safe unwrap since physical indices came from a valid run array.
1504    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    // SAFETY: run-ends are strictly increasing with last value == logical length.
1510    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    // SAFETY: `new_values` has one entry per run.
1521    Ok(
1522        unsafe {
1523            RunArray::<T>::new_unchecked(run_array.data_type().clone(), run_ends, new_values)
1524        },
1525    )
1526}
1527
1528/// Physical run index for each logical take slot.
1529///
1530/// `None` means the logical index is null. Only valid indices are passed to
1531/// [`RunArray::get_physical_indices`]; a null slot's backing integer is ignored
1532/// and may be out of range.
1533fn 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
1579/// Takes/filters a fixed size list array's inner data using the offsets of the list array.
1580fn 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            // Safety: Range always has known length.
1601            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
1612/// To avoid generating take implementations for every index type, instead we
1613/// only generate for UInt32 and UInt64 and coerce inputs to these types
1614trait 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
1670/// Take rows by index from [`RecordBatch`] and returns a new [`RecordBatch`] from those indexes.
1671///
1672/// This function will call [`take`] on each array of the [`RecordBatch`] and assemble a new [`RecordBatch`].
1673///
1674/// # Example
1675/// ```
1676/// # use std::sync::Arc;
1677/// # use arrow_array::{StringArray, Int32Array, UInt32Array, RecordBatch};
1678/// # use arrow_schema::{DataType, Field, Schema};
1679/// # use arrow_select::take::take_record_batch;
1680/// let schema = Arc::new(Schema::new(vec![
1681///     Field::new("a", DataType::Int32, true),
1682///     Field::new("b", DataType::Utf8, true),
1683/// ]));
1684/// let batch = RecordBatch::try_new(
1685///     schema.clone(),
1686///     vec![
1687///         Arc::new(Int32Array::from_iter_values(0..20)),
1688///         Arc::new(StringArray::from_iter_values(
1689///             (0..20).map(|i| format!("str-{}", i)),
1690///         )),
1691///     ],
1692/// )
1693/// .unwrap();
1694///
1695/// let indices = UInt32Array::from(vec![1, 5, 10]);
1696/// let taken = take_record_batch(&batch, &indices).unwrap();
1697///
1698/// let expected = RecordBatch::try_new(
1699///     schema,
1700///     vec![
1701///         Arc::new(Int32Array::from(vec![1, 5, 10])),
1702///         Arc::new(StringArray::from(vec!["str-1", "str-5", "str-10"])),
1703///     ],
1704/// )
1705/// .unwrap();
1706/// assert_eq!(taken, expected);
1707/// ```
1708pub 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            // Safety: indices were validated by the first take_impl call
1724            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    // create a simple struct for testing purposes
1829    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        // int8
1968        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        // int16
1977        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        // int32
1986        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        // int64
1995        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        // uint8
2004        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        // uint16
2013        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        // uint32
2022        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        // int64
2031        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        // interval_year_month
2040        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        // interval_day_time
2049        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        // interval_month_day_nano
2061        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        // duration_second
2073        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        // duration_millisecond
2082        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        // duration_microsecond
2091        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        // duration_nanosecond
2100        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        // float32
2109        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        // float64
2118        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        // int16
2150        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        // int64
2158        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        // uint64
2166        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        // duration_millisecond
2174        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        // float32
2182        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        // int16
2195        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        // duration_millisecond
2203        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        // float32
2211        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        // boolean
2223        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        // indices where the masked invalid elements would be out of bounds
2234        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        // indices where the masked invalid elements would be out of bounds
2251        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        // boolean
2275        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    // Null indices + sliced source boolean array — exercises src_offset in the sparse take_bits path.
2285    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); // logical: [false, true, false], offset=1
2288        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    // >8 elements: exercises the 8-at-a-time unrolled byte packing in take_bits.
2300    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    // Sliced source: verifies src_offset is applied when no null indices.
2315    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); // [true, false, true], offset=2
2318        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    // Nullable source, >8 elements: exercises the 8-at-a-time unrolled byte packing in take_bits_with_validity.
2328    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    // Nullable source, null indices, sliced source: verifies src_offset with both null paths.
2361    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); // [false, null, true, false], offset=1
2365        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    /// Take from a *sliced* byte array, i.e. one whose value offsets do not
2425    /// start at zero. This exercises copying byte data out of an array with a
2426    /// non-zero base offset for both the no-null fast path and the nullable
2427    /// path (null indices and selected null values).
2428    #[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        // Slice so the underlying value offsets no longer start at 0:
2440        // sliced == [None, "ccccc", "dd", None, "eeee"]
2441        let sliced = values.slice(2, 5);
2442
2443        // Fast path: every output slot is valid (no null indices, no null
2444        // values selected).
2445        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        // Nullable path: a null index (position 1) and selected null values
2452        // (sliced indices 0 and 3 are null).
2453        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            // ["hello", "world", null, "large payload over 12 bytes", "lulu"]
2468            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            // ["large payload over 12 bytes", null, "world", "large payload over 12 bytes", "lulu", null]
2486            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            // Construct a value array, [[0,0,0], [-1,-2,-1], [], [2,3]]
2512            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            // index returns: [[2,3], null, [-1,-2,-1], [], [0,0,0]]
2519            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            // construct a value array with expected results:
2525            // [[2,3], null, [-1,-2,-1], [], [0,0,0]]
2526            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                // null buffer remains the same as only the indices have nulls
2543                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            // Construct a value array, [[0,null,0], [-1,-2,3], [null], [5,null]]
2553            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            // index returns: [[null], null, [-1,-2,3], [2,null], [0,null,0]]
2570            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            // construct a value array with expected results:
2576            // [[null], null, [-1,-2,3], [5,null], [0,null,0]]
2577            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                // null buffer remains the same as only the indices have nulls
2595                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            // Construct a value array, [[0,null,0], [-1,-2,3], null, [5,null]]
2605            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            // the entry at index 2 is null
2619            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            // index returns: [null, null, [-1,-2,3], [5,null], [0,null,0]]
2624            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            // construct a value array with expected results:
2630            // [null, null, [-1,-2,3], [5,null], [0,null,0]]
2631            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            // null buffer must be recalculated as both values and indices have nulls
2644            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    // Fast path (primitive child, no child nulls) with null indices — verifies offset backfill for null slots.
2728    fn test_take_list_primitive_child_null_indices() {
2729        // Row sizes deliberately vary (1, 3, 2) so the test exercises
2730        // non-uniform offset arithmetic, not just uniform stride.
2731        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        // Take reversed indices
2782        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        // Take with null indices
2800        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        // Take at null values
2814        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        // Take null indices/values, with slicing.
2828        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        // The two middle indices are null -> Should be null in the output.
2916        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    /// The [`take_fixed_size_binary`] kernel contains optimizations that provide a faster
2928    /// implementation for commonly-used value lengths. This test uses a value length that is not
2929    /// optimized to test both code paths.
2930    #[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        // The two middle indices are null -> Should be null in the output.
2945        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        // Construct a value array, [[0,0,0], [-1,-2,-1], [2,3]]
2960        let values = Arc::new(Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2, 3]));
2961        // Construct offsets
2962        let value_offsets = OffsetBuffer::<i32>::new(ScalarBuffer::from(vec![0, 3, 6, 8]));
2963        // Construct a list array from the above two
2964        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        // A panic is expected here since we have not supplied the check_bounds
2970        // option.
2971        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()); // 2 because of indices, 1 because of struct array
3047
3048        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        // int64
3066        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), // first "foo"
3140            Some(7), // last "foo"
3141            None,    // null index should return null
3142            Some(5), // second "bar"
3143            Some(6), // another "bar"
3144            Some(2), // empty string
3145            Some(3), // input is null at this index
3146        ]);
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        // dictionary values should stay the same
3157        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        // Takes the same rows from a two-row slice of a large array and from an unsliced
3260        // copy of those two rows. The results should be equal and use the same memory.
3261        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        // Lists of 10 values per row
3274        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        // adjacent identical values are merged: [2,2,2,2,2,1,1] -> 2 runs
3304        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        // A null index must not become logical index 0, and null indices must merge with
3318        // consecutive runs whose values are already null.
3319        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); // [3, 3, 3, 4, 4, 5]
3361
3362        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        // [3, 5, 5, 3, 4] -> 4 physical runs (no adjacent duplicates to merge)
3368        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        // Build indices with values that are out of bounds, but masked by null mask
3410        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        // Dense [1, 2, 3]; a null index must not become child offset 0. Use a non-zero
3565        // type id to verify null indices don't use the invalid default type id 0.
3566        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        // Use an out-of-bounds backing value for the null index to exercise `take_native`'s
3584        // default type id path. A union need not have a child with type id 0.
3585        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        // The inner union field is marked nullable, but it has no nullable
3817        // children, so it cannot represent a take-null. The sibling Int32 field
3818        // can, and must be selected instead of erroring.
3819        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        // A nullable-marked nested union is not a valid null child when none of
3859        // its own children can store a null.
3860        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        // A nullable-marked REE child cannot represent a take-null when its
3900        // values field is non-nullable. The sibling Int32 field must be used.
3901        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    /// Fixture for the offset-overflow tests: a single large value plus the
3985    /// number of times it must be selected so the cumulative offset exceeds
3986    /// `i32::MAX`. Using a large value keeps the index count (and the test
3987    /// runtime) small.
3988    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    /// The offset-overflow error must also be produced on the nullable code
4006    /// path (when the output contains nulls), not only on the no-null fast path.
4007    #[test]
4008    fn test_take_bytes_offset_overflow_nullable() {
4009        let (values, n) = offset_overflow_fixture();
4010        // A null index forces the output to contain nulls, exercising the
4011        // nullable code path.
4012        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        // Verify the result is a valid empty RunArray
4034        assert_eq!(result.len(), 0);
4035        assert_eq!(result.null_count(), 0);
4036
4037        // Verify that the result can be downcast and used without validation errors
4038        // This specifically tests that "The values in run_ends array should be strictly positive" is not triggered
4039        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        // https://github.com/apache/arrow-rs/issues/7710
4050        // Indices [0,1,4,5] select from [1,1,0,0,1,1] — the 0s are skipped,
4051        // so the output should be a single run of 1s, not two.
4052        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        // Verify physical layout: all four logical values collapse into one run.
4064        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        // Verify physical layout: all four logical values collapse into one run.
4093        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        // Validates that runs are merged whether the same logical value comes
4107        // from the same physical index (repeated indices) or distinct physical indices.
4108        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        // [bob,bob,bob,bob,bob,alice,alice,alice,eve,eve,eve]
4117        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        // Verify physical layout: 11 logical values across exactly 3 physical runs.
4125
4126        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    // parent null rows on FixedSizeList must propagate to child elements.
4146    #[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}