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/// Verifies that the non-null values of `indices` are all `< len`
167fn check_bounds<T: ArrowPrimitiveType>(
168    len: usize,
169    indices: &PrimitiveArray<T>,
170) -> Result<(), ArrowError>
171where
172    T::Native: Display,
173{
174    let Some(len) = T::Native::from_usize(len) else {
175        return if T::DATA_TYPE.is_integer() {
176            // the biggest representable value for T::Native is lower than len, e.g: u8::MAX < 512, no need to check bounds
177            Ok(())
178        } else {
179            Err(ArrowError::ComputeError("Cast to usize failed".to_string()))
180        };
181    };
182
183    if indices.null_count() > 0 {
184        indices.iter().flatten().try_for_each(|index| {
185            if index >= len {
186                return Err(ArrowError::ComputeError(format!(
187                    "Array index out of bounds, cannot get item at index {index} from {len} entries"
188                )));
189            }
190            Ok(())
191        })
192    } else {
193        let in_bounds = indices.values().iter().fold(true, |in_bounds, &i| {
194            in_bounds & (i >= T::Native::ZERO) & (i < len)
195        });
196
197        if !in_bounds {
198            for &index in indices.values() {
199                if index < T::Native::ZERO || index >= len {
200                    return Err(ArrowError::ComputeError(format!(
201                        "Array index out of bounds, cannot get item at index {index} from {len} entries"
202                    )));
203                }
204            }
205        }
206
207        Ok(())
208    }
209}
210
211#[inline(never)]
212fn take_impl<IndexType: ArrowPrimitiveType, const CHECKED: bool>(
213    values: &dyn Array,
214    indices: &PrimitiveArray<IndexType>,
215) -> Result<ArrayRef, ArrowError> {
216    if indices.is_empty() {
217        if let DataType::Union(fields, _) = values.data_type()
218            && fields.is_empty()
219        {
220            // `new_empty_array` cannot construct a union with no fields, but an existing empty
221            // union can be sliced without materializing a child array.
222            return Ok(values.slice(0, 0));
223        }
224        return Ok(new_empty_array(values.data_type()));
225    }
226    downcast_primitive_array! {
227        values => Ok(Arc::new(take_primitive::<_, _, CHECKED>(values, indices)?)),
228        DataType::Boolean => {
229            let values = values.as_any().downcast_ref::<BooleanArray>().unwrap();
230            Ok(Arc::new(take_boolean::<_, CHECKED>(values, indices)))
231        }
232        DataType::Utf8 => {
233            Ok(Arc::new(take_bytes::<_, _, CHECKED>(values.as_string::<i32>(), indices)?))
234        }
235        DataType::LargeUtf8 => {
236            Ok(Arc::new(take_bytes::<_, _, CHECKED>(values.as_string::<i64>(), indices)?))
237        }
238        DataType::Utf8View => {
239            Ok(Arc::new(take_byte_view::<_, _, CHECKED>(values.as_string_view(), indices)?))
240        }
241        DataType::List(_) => {
242            Ok(Arc::new(take_list::<_, Int32Type, CHECKED>(values.as_list(), indices)?))
243        }
244        DataType::LargeList(_) => {
245            Ok(Arc::new(take_list::<_, Int64Type, CHECKED>(values.as_list(), indices)?))
246        }
247        DataType::ListView(_) => {
248            Ok(Arc::new(take_list_view::<_, Int32Type, CHECKED>(values.as_list_view(), indices)?))
249        }
250        DataType::LargeListView(_) => {
251            Ok(Arc::new(take_list_view::<_, Int64Type, CHECKED>(values.as_list_view(), indices)?))
252        }
253        DataType::FixedSizeList(_, length) => {
254            let values = values
255                .as_any()
256                .downcast_ref::<FixedSizeListArray>()
257                .unwrap();
258            Ok(Arc::new(take_fixed_size_list::<_, CHECKED>(
259                values,
260                indices,
261                *length as u32,
262            )?))
263        }
264        DataType::Map(field, ordered) => {
265            let list_arr = ListArray::from(values.as_map().clone());
266            let list_data = take_list::<_, Int32Type, CHECKED>(&list_arr, indices)?;
267            let (_, offsets, entries, nulls) = list_data.into_parts();
268            let entries = entries.as_struct().clone();
269            Ok(Arc::new(MapArray::try_new(
270                field.clone(),
271                offsets,
272                entries,
273                nulls,
274                *ordered,
275            )?))
276        }
277        DataType::Struct(fields) => {
278            let array: &StructArray = values.as_struct();
279            let arrays  = array
280                .columns()
281                .iter()
282                .map(|a| take_impl::<_, CHECKED>(a.as_ref(), indices))
283                .collect::<Result<Vec<ArrayRef>, _>>()?;
284            let fields: Vec<(FieldRef, ArrayRef)> =
285                fields.iter().cloned().zip(arrays).collect();
286
287            // Create the null bit buffer.
288            let is_valid: Buffer = indices
289                .iter()
290                .map(|index| {
291                    if let Some(index) = index {
292                        array.is_valid(index.to_usize().unwrap())
293                    } else {
294                        false
295                    }
296                })
297                .collect();
298
299            if fields.is_empty() {
300                let nulls = NullBuffer::new(BooleanBuffer::new(is_valid, 0, indices.len()));
301                Ok(Arc::new(StructArray::new_empty_fields(indices.len(), Some(nulls))))
302            } else {
303                Ok(Arc::new(StructArray::from((fields, is_valid))) as ArrayRef)
304            }
305        }
306        DataType::Dictionary(_, _) => downcast_dictionary_array! {
307            values => Ok(Arc::new(take_dict::<_, _, CHECKED>(values, indices)?)),
308            t => unimplemented!("Take not supported for dictionary type {:?}", t)
309        }
310        DataType::RunEndEncoded(_, _) => downcast_run_array! {
311            values => Ok(Arc::new(take_run(values, indices)?)),
312            t => unimplemented!("Take not supported for run type {:?}", t)
313        }
314        DataType::Binary => {
315            Ok(Arc::new(take_bytes::<_, _, CHECKED>(values.as_binary::<i32>(), indices)?))
316        }
317        DataType::LargeBinary => {
318            Ok(Arc::new(take_bytes::<_, _, CHECKED>(values.as_binary::<i64>(), indices)?))
319        }
320        DataType::BinaryView => {
321            Ok(Arc::new(take_byte_view::<_, _, CHECKED>(values.as_binary_view(), indices)?))
322        }
323        DataType::FixedSizeBinary(size) => {
324            let values = values
325                .as_any()
326                .downcast_ref::<FixedSizeBinaryArray>()
327                .unwrap();
328            Ok(Arc::new(take_fixed_size_binary::<_, CHECKED>(values, indices, *size)?))
329        }
330        DataType::Null => {
331            // Take applied to a null array produces a null array.
332            if values.len() >= indices.len() {
333                // If the existing null array is as big as the indices, we can use a slice of it
334                // to avoid allocating a new null array.
335                Ok(values.slice(0, indices.len()))
336            } else {
337                // If the existing null array isn't big enough, create a new one.
338                Ok(new_null_array(&DataType::Null, indices.len()))
339            }
340        }
341        DataType::Union(fields, UnionMode::Sparse) => {
342            let mut children = Vec::with_capacity(fields.len());
343            let values = values.as_any().downcast_ref::<UnionArray>().unwrap();
344            let type_ids = take_union_type_ids(fields, values.type_ids(), indices)?;
345            for (type_id, _field) in fields.iter() {
346                let values = values.child(type_id);
347                let values = take_impl::<_, CHECKED>(values, indices)?;
348                children.push(values);
349            }
350            let array = UnionArray::try_new(fields.clone(), type_ids, None, children)?;
351            Ok(Arc::new(array))
352        }
353        DataType::Union(fields, UnionMode::Dense) => {
354            let values = values.as_any().downcast_ref::<UnionArray>().unwrap();
355
356            let type_ids = PrimitiveArray::<Int8Type>::try_new(
357                take_union_type_ids(fields, values.type_ids(), indices)?,
358                None,
359            )?;
360            // Keep index nulls so `take` of each child writes a null instead of
361            // reading child offset 0 (the default `take_native` fills in).
362            let offsets = <PrimitiveArray<Int32Type>>::try_new(
363                take_native(values.offsets().unwrap(), indices),
364                indices.nulls().cloned(),
365            )?;
366
367            let children = fields.iter()
368                .map(|(field_type_id, _)| {
369                    let mask = BooleanArray::from_unary(&type_ids, |value_type_id| value_type_id == field_type_id);
370
371                    let indices = crate::filter::filter(&offsets, &mask)?;
372
373                    let values = values.child(field_type_id);
374
375                    take_impl::<_, CHECKED>(values, indices.as_primitive::<Int32Type>())
376                })
377                .collect::<Result<_, _>>()?;
378
379            let mut child_offsets = [0; 128];
380
381            let offsets = type_ids.values()
382                .iter()
383                .map(|&i| {
384                    let offset = child_offsets[i as usize];
385
386                    child_offsets[i as usize] += 1;
387
388                    offset
389                })
390                .collect();
391
392            let (_, type_ids, _) = type_ids.into_parts();
393
394            let array = UnionArray::try_new(fields.clone(), type_ids, Some(offsets), children)?;
395
396            Ok(Arc::new(array))
397        }
398        t => unimplemented!("Take not supported for data type {:?}", t)
399    }
400}
401
402/// Takes union type ids, substituting a valid type id for null take indices.
403///
404/// Union arrays do not have a top-level null bitmap. A null is represented by selecting an
405/// arbitrary valid child type id with a null value in that child. In particular, a null index
406/// cannot fall back to type id `0`, as unions are not required to have such a child.
407fn take_union_type_ids<IndexType: ArrowPrimitiveType>(
408    fields: &UnionFields,
409    type_ids: &ScalarBuffer<i8>,
410    indices: &PrimitiveArray<IndexType>,
411) -> Result<ScalarBuffer<i8>, ArrowError> {
412    if indices.null_count() == 0 {
413        return Ok(take_native(type_ids, indices));
414    }
415
416    let null_type_id = fields
417        .iter()
418        .next()
419        .map(|(type_id, _)| type_id)
420        .ok_or_else(|| {
421            ArrowError::ComputeError(
422                "Cannot take from a union with zero fields when indices contains nulls".into(),
423            )
424        })?;
425    let taken_type_ids = take_native(type_ids, indices);
426    let type_ids = indices
427        .iter()
428        .zip(&taken_type_ids)
429        .map(|(index, &type_id)| {
430            if index.is_some() {
431                type_id
432            } else {
433                null_type_id
434            }
435        })
436        .collect::<ScalarBuffer<_>>();
437    Ok(type_ids)
438}
439
440/// Options that define how `take` should behave
441#[derive(Clone, Debug, Default)]
442pub struct TakeOptions {
443    /// Perform bounds check before taking indices from values.
444    /// If enabled, an `ArrowError` is returned if the indices are out of bounds.
445    /// If not enabled, and indices exceed bounds, the kernel will panic.
446    pub check_bounds: bool,
447}
448
449/// `take` implementation for all primitive arrays
450///
451/// This checks if an `indices` slot is populated, and gets the value from `values`
452///  as the populated index.
453/// If the `indices` slot is null, a null value is returned.
454/// For example, given:
455///     values:  [1, 2, 3, null, 5]
456///     indices: [0, null, 4, 3]
457/// The result is: [1 (slot 0), null (null slot), 5 (slot 4), null (slot 3)]
458fn take_primitive<T, I, const CHECKED: bool>(
459    values: &PrimitiveArray<T>,
460    indices: &PrimitiveArray<I>,
461) -> Result<PrimitiveArray<T>, ArrowError>
462where
463    T: ArrowPrimitiveType,
464    I: ArrowPrimitiveType,
465{
466    let values_buf = take_native(values.values(), indices);
467    let nulls = take_nulls::<_, CHECKED>(values.nulls(), indices);
468    Ok(PrimitiveArray::try_new(values_buf, nulls)?.with_data_type(values.data_type().clone()))
469}
470
471#[inline(never)]
472fn take_nulls<I: ArrowPrimitiveType, const CHECKED: bool>(
473    values: Option<&NullBuffer>,
474    indices: &PrimitiveArray<I>,
475) -> Option<NullBuffer> {
476    match values.filter(|n| n.null_count() > 0) {
477        Some(n) => NullBuffer::from_unsliced_buffer(
478            take_bits::<_, CHECKED>(n.inner(), indices).into_inner(),
479            indices.len(),
480        ),
481        None => indices.nulls().cloned(),
482    }
483}
484
485#[inline(never)]
486fn take_native<T: ArrowNativeType, I: ArrowPrimitiveType>(
487    values: &[T],
488    indices: &PrimitiveArray<I>,
489) -> ScalarBuffer<T> {
490    match indices.nulls().filter(|n| n.null_count() > 0) {
491        Some(n) => indices
492            .values()
493            .iter()
494            .enumerate()
495            .map(|(idx, index)| match values.get(index.as_usize()) {
496                Some(v) => *v,
497                // SAFETY: idx<indices.len()
498                None => match unsafe { n.inner().value_unchecked(idx) } {
499                    false => T::default(),
500                    true => panic!("Out-of-bounds index {index:?}"),
501                },
502            })
503            .collect(),
504        None => indices
505            .values()
506            .iter()
507            .map(|index| values[index.as_usize()])
508            .collect(),
509    }
510}
511
512/// Read the bit at `src_bit_idx` from `src` and, if it is set, write a `1` to `dst_bit_idx`
513/// in `dst`. Leaves `dst_bit_idx` unchanged (zero) when the source bit is unset.
514///
515/// ```text
516/// src = 0b00100000  (bit 5 is set)
517/// copy_bit_if_set(src, 5, dst, 2)  →  dst bit 2 becomes 1
518/// ```
519///
520/// # Safety
521/// - `src` must be valid for reads up to byte `src_bit_idx / 8`.
522/// - `dst` must be valid for writes up to byte `dst_bit_idx / 8`.
523#[inline(always)]
524unsafe fn copy_bit_if_set(src: *const u8, src_bit_idx: usize, dst: *mut u8, dst_bit_idx: usize) {
525    unsafe {
526        if bit_util::get_bit_raw(src, src_bit_idx) {
527            bit_util::set_bit_raw(dst, dst_bit_idx);
528        }
529    }
530}
531
532/// Read the bit at `bit_idx` from `src` and return it shifted to `out_pos`, ready to be
533/// OR'd into an output byte accumulator.
534///
535/// ```text
536/// src = 0b10100000  (bit 5 is set)
537/// pack_bit(src, 5, 2)  →  0b00000100   (bit from position 5, placed at position 2)
538/// ```
539///
540/// # Safety
541/// `src` must be valid for reads up to byte `bit_idx / 8`.
542#[inline(always)]
543unsafe fn pack_bit(src: *const u8, bit_idx: usize, out_pos: usize) -> u8 {
544    let byte = unsafe { *src.add(bit_idx >> 3) }; // byte containing bit `bit_idx`
545    ((byte >> (bit_idx & 7)) & 1) << out_pos // extract the bit, shift to output position
546}
547
548#[inline(never)]
549fn take_bits<I: ArrowPrimitiveType, const CHECKED: bool>(
550    values: &BooleanBuffer,
551    indices: &PrimitiveArray<I>,
552) -> BooleanBuffer {
553    let len = indices.len();
554    let src_offset = values.offset();
555    let src_ptr = values.values().as_ptr();
556    let out_bytes = len.div_ceil(8);
557
558    match indices.nulls().filter(|nulls| nulls.null_count() > 0) {
559        Some(index_nulls) => {
560            let mut output = vec![0u8; out_bytes];
561            let out_ptr = output.as_mut_ptr();
562            index_nulls.valid_indices().for_each(|valid_idx| {
563                // SAFETY: valid_idx < indices.len(), guaranteed by valid_indices().
564                let index_val = unsafe { indices.value_unchecked(valid_idx) }.as_usize();
565                if CHECKED {
566                    if values.value(index_val) {
567                        // SAFETY: valid_idx < indices.len() = len, output buffer holds len bits.
568                        unsafe { bit_util::set_bit_raw(out_ptr, valid_idx) };
569                    }
570                } else {
571                    // SAFETY: caller guarantees index_val < values.len().
572                    unsafe { copy_bit_if_set(src_ptr, index_val + src_offset, out_ptr, valid_idx) };
573                }
574            });
575            BooleanBuffer::new(Buffer::from(output), 0, len)
576        }
577        None => {
578            // Build the output byte-by-byte with an 8-element inner loop so the
579            // compiler can fully unroll it and issue the 8 source loads in parallel.
580            let mut output = vec![0u8; out_bytes];
581            let out_slice = output.as_mut_slice();
582            let full_bytes = len / 8;
583
584            for (byte_idx, out_byte) in out_slice.iter_mut().enumerate().take(full_bytes) {
585                let base = byte_idx * 8;
586                let mut byte = 0u8;
587                for bit in 0..8usize {
588                    // SAFETY: base + bit < full_bytes * 8 <= len, so base + bit is a valid
589                    // position in the indices array.
590                    let index_val = unsafe { indices.value_unchecked(base + bit) }.as_usize();
591                    if CHECKED {
592                        byte |= (values.value(index_val) as u8) << bit;
593                    } else {
594                        // SAFETY: caller guarantees index_val < values.len().
595                        byte |= unsafe { pack_bit(src_ptr, index_val + src_offset, bit) };
596                    }
597                }
598                *out_byte = byte;
599            }
600            // Handle remaining bits when len is not a multiple of 8.
601            if full_bytes < out_bytes {
602                let base = full_bytes * 8;
603                let mut byte = 0u8;
604                for bit in 0..(len - base) {
605                    // SAFETY: base + bit < len (remainder loop bound), so base + bit is a
606                    // valid position in the indices array.
607                    let index_val = unsafe { indices.value_unchecked(base + bit) }.as_usize();
608                    if CHECKED {
609                        byte |= (values.value(index_val) as u8) << bit;
610                    } else {
611                        // SAFETY: caller guarantees index_val < values.len().
612                        byte |= unsafe { pack_bit(src_ptr, index_val + src_offset, bit) };
613                    }
614                }
615                out_slice[full_bytes] = byte;
616            }
617            BooleanBuffer::new(Buffer::from(output), 0, len)
618        }
619    }
620}
621
622/// Gather value bits and validity bits from two boolean buffers in a single pass.
623/// Used when the values array itself has nulls, avoiding two separate `take_bits` calls.
624#[inline(never)]
625fn take_bits_with_validity<I: ArrowPrimitiveType, const CHECKED: bool>(
626    values: &BooleanBuffer,
627    validity: &BooleanBuffer,
628    indices: &PrimitiveArray<I>,
629) -> (BooleanBuffer, Option<NullBuffer>) {
630    let len = indices.len();
631    let value_bit_offset = values.offset();
632    let validity_bit_offset = validity.offset();
633    let value_data_ptr = values.values().as_ptr();
634    let validity_data_ptr = validity.values().as_ptr();
635    let out_bytes = len.div_ceil(8);
636
637    let mut value_out = vec![0u8; out_bytes];
638    let mut validity_out = vec![0u8; out_bytes];
639
640    match indices.nulls().filter(|nulls| nulls.null_count() > 0) {
641        Some(index_nulls) => {
642            let value_out_ptr = value_out.as_mut_ptr();
643            let validity_out_ptr = validity_out.as_mut_ptr();
644            for out_pos in index_nulls.valid_indices() {
645                // SAFETY: out_pos < indices.len(), guaranteed by valid_indices().
646                let src_idx = unsafe { indices.value_unchecked(out_pos) }.as_usize();
647                if CHECKED {
648                    if values.value(src_idx) {
649                        // SAFETY: out_pos < indices.len() = len, output buffer holds len bits.
650                        unsafe { bit_util::set_bit_raw(value_out_ptr, out_pos) };
651                    }
652                    if validity.value(src_idx) {
653                        // SAFETY: out_pos < indices.len() = len, output buffer holds len bits.
654                        unsafe { bit_util::set_bit_raw(validity_out_ptr, out_pos) };
655                    }
656                } else {
657                    // SAFETY: caller guarantees src_idx < values.len().
658                    unsafe {
659                        copy_bit_if_set(
660                            value_data_ptr,
661                            src_idx + value_bit_offset,
662                            value_out_ptr,
663                            out_pos,
664                        );
665                        copy_bit_if_set(
666                            validity_data_ptr,
667                            src_idx + validity_bit_offset,
668                            validity_out_ptr,
669                            out_pos,
670                        );
671                    }
672                }
673            }
674        }
675        None => {
676            let value_out_slice = value_out.as_mut_slice();
677            let validity_out_slice = validity_out.as_mut_slice();
678            let full_bytes = len / 8;
679
680            for (byte_idx, (value_out_byte, validity_out_byte)) in value_out_slice
681                .iter_mut()
682                .zip(validity_out_slice.iter_mut())
683                .enumerate()
684                .take(full_bytes)
685            {
686                let bit_base = byte_idx * 8;
687                let mut packed_values = 0u8;
688                let mut packed_validity = 0u8;
689                for bit_pos in 0..8usize {
690                    // SAFETY: bit_base + bit_pos < full_bytes * 8 <= len.
691                    let src_idx = unsafe { indices.value_unchecked(bit_base + bit_pos) }.as_usize();
692                    if CHECKED {
693                        packed_values |= (values.value(src_idx) as u8) << bit_pos;
694                        packed_validity |= (validity.value(src_idx) as u8) << bit_pos;
695                    } else {
696                        // SAFETY: caller guarantees src_idx < values.len().
697                        packed_values |= unsafe {
698                            pack_bit(value_data_ptr, src_idx + value_bit_offset, bit_pos)
699                        };
700                        packed_validity |= unsafe {
701                            pack_bit(validity_data_ptr, src_idx + validity_bit_offset, bit_pos)
702                        };
703                    }
704                }
705                *value_out_byte = packed_values;
706                *validity_out_byte = packed_validity;
707            }
708            // Handle remaining bits when len is not a multiple of 8.
709            if full_bytes < out_bytes {
710                let bit_base = full_bytes * 8;
711                let mut packed_values = 0u8;
712                let mut packed_validity = 0u8;
713                for bit_pos in 0..(len - bit_base) {
714                    // SAFETY: bit_base + bit_pos < len (remainder loop bound).
715                    let src_idx = unsafe { indices.value_unchecked(bit_base + bit_pos) }.as_usize();
716                    if CHECKED {
717                        packed_values |= (values.value(src_idx) as u8) << bit_pos;
718                        packed_validity |= (validity.value(src_idx) as u8) << bit_pos;
719                    } else {
720                        // SAFETY: caller guarantees src_idx < values.len().
721                        packed_values |= unsafe {
722                            pack_bit(value_data_ptr, src_idx + value_bit_offset, bit_pos)
723                        };
724                        packed_validity |= unsafe {
725                            pack_bit(validity_data_ptr, src_idx + validity_bit_offset, bit_pos)
726                        };
727                    }
728                }
729                value_out_slice[full_bytes] = packed_values;
730                validity_out_slice[full_bytes] = packed_validity;
731            }
732        }
733    }
734
735    let value_buf_out = BooleanBuffer::new(Buffer::from(value_out), 0, len);
736    let validity_buf_out = NullBuffer::from_unsliced_buffer(validity_out, len);
737    (value_buf_out, validity_buf_out)
738}
739
740/// `take` implementation for boolean arrays
741fn take_boolean<IndexType: ArrowPrimitiveType, const CHECKED: bool>(
742    array: &BooleanArray,
743    indices: &PrimitiveArray<IndexType>,
744) -> BooleanArray {
745    let bits = array.values();
746    match array.nulls().filter(|n| n.null_count() > 0) {
747        Some(array_nulls) => {
748            let (val_buf, null_buf) =
749                take_bits_with_validity::<_, CHECKED>(bits, array_nulls.inner(), indices);
750            BooleanArray::new(val_buf, null_buf)
751        }
752        None => {
753            let val_buf = take_bits::<_, CHECKED>(bits, indices);
754            let null_buf = take_nulls::<_, CHECKED>(None, indices);
755            BooleanArray::new(val_buf, null_buf)
756        }
757    }
758}
759
760/// `take` implementation for string arrays
761fn take_bytes<T: ByteArrayType, IndexType: ArrowPrimitiveType, const CHECKED: bool>(
762    array: &GenericByteArray<T>,
763    indices: &PrimitiveArray<IndexType>,
764) -> Result<GenericByteArray<T>, ArrowError> {
765    let mut values: Vec<u8> = Vec::new();
766    let mut offsets = Vec::with_capacity(indices.len() + 1);
767    offsets.push(T::Offset::default());
768
769    let input_offsets = array.value_offsets();
770    let mut capacity = 0;
771    let nulls = take_nulls::<_, CHECKED>(array.nulls(), indices);
772
773    // Branch on output nulls — `None` means every output slot is valid.
774    match nulls.as_ref().filter(|n| n.null_count() > 0) {
775        // Fast path: no nulls in output, every index is valid.
776        None => {
777            for index in indices.values() {
778                let index = index.as_usize();
779                let start = input_offsets[index].as_usize();
780                let end = input_offsets[index + 1].as_usize();
781                capacity += end - start;
782                offsets.push(
783                    T::Offset::from_usize(capacity)
784                        .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?,
785                );
786            }
787
788            values.reserve(capacity);
789
790            let dst = values.spare_capacity_mut();
791            debug_assert!(dst.len() >= capacity);
792            let mut offset = 0;
793
794            for index in indices.values() {
795                // SAFETY: in-bounds proven by the first loop's bounds-checked offset access.
796                // dst asserted above to include the required capacity.
797                unsafe {
798                    let data: &[u8] = array.value_unchecked(index.as_usize()).as_ref();
799                    std::ptr::copy_nonoverlapping(
800                        data.as_ptr(),
801                        dst.get_unchecked_mut(offset..).as_mut_ptr().cast::<u8>(),
802                        data.len(),
803                    );
804                    offset += data.len();
805                }
806            }
807
808            // SAFETY: wrote exactly `capacity` bytes above; reserved on line above.
809            unsafe {
810                values.set_len(capacity);
811            }
812        }
813        // Nullable path: only process valid (non-null) output positions.
814        Some(output_nulls) => {
815            let mut source_ranges = Vec::with_capacity(indices.len() - output_nulls.null_count());
816            let mut last_filled = 0;
817
818            // Pre-fill offsets; we overwrite valid positions below.
819            offsets.resize(indices.len() + 1, T::Offset::default());
820
821            // Pass 1: find all valid ranges that need to be copied.
822            for i in output_nulls.valid_indices() {
823                let current_offset = T::Offset::from_usize(capacity)
824                    .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?;
825                // Fill offsets for skipped null slots so they get zero-length ranges.
826                if last_filled < i {
827                    offsets[last_filled + 1..=i].fill(current_offset);
828                }
829
830                // SAFETY: `i` comes from a validity bitmap over `indices`, so it is in-bounds.
831                let index = unsafe { indices.value_unchecked(i) }.as_usize();
832                let start = input_offsets[index].as_usize();
833                let end = input_offsets[index + 1].as_usize();
834                capacity += end - start;
835                offsets[i + 1] = T::Offset::from_usize(capacity)
836                    .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?;
837
838                source_ranges.push((start, end));
839                last_filled = i + 1;
840            }
841
842            // Fill trailing null offsets after the last valid position.
843            let final_offset = T::Offset::from_usize(capacity)
844                .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?;
845            offsets[last_filled + 1..].fill(final_offset);
846            // Pass 2: copy byte data for all collected ranges.
847            values.reserve(capacity);
848            debug_assert_eq!(
849                source_ranges.iter().map(|(s, e)| e - s).sum::<usize>(),
850                capacity,
851                "capacity must equal total bytes across all ranges"
852            );
853
854            let src = array.value_data();
855            let src = src.as_ptr();
856            let dst = values.spare_capacity_mut();
857            debug_assert!(dst.len() >= capacity);
858
859            let mut offset = 0;
860
861            for (start, end) in source_ranges {
862                let value_len = end - start;
863                // SAFETY: caller guarantees each (start, end) is in-bounds of `src`.
864                // `dst` asserted above to include the required capacity.
865                // The regions don't overlap (src is input, dst is a fresh allocation).
866                unsafe {
867                    std::ptr::copy_nonoverlapping(
868                        src.add(start),
869                        dst.get_unchecked_mut(offset..).as_mut_ptr().cast::<u8>(),
870                        value_len,
871                    );
872                    offset += value_len;
873                }
874            }
875            // SAFETY: caller guarantees `capacity` == total bytes across all ranges,
876            // so the loop above wrote exactly `capacity` bytes.
877            unsafe { values.set_len(capacity) };
878        }
879    }
880
881    // SAFETY: offsets are monotonically increasing and in-bounds of `values`,
882    // and `nulls` (if present) has length == `indices.len()`.
883    let array = unsafe {
884        let offsets = OffsetBuffer::new_unchecked(offsets.into());
885        GenericByteArray::<T>::new_unchecked(offsets, values.into(), nulls)
886    };
887
888    Ok(array)
889}
890
891/// `take` implementation for byte view arrays
892fn take_byte_view<T: ByteViewType, IndexType: ArrowPrimitiveType, const CHECKED: bool>(
893    array: &GenericByteViewArray<T>,
894    indices: &PrimitiveArray<IndexType>,
895) -> Result<GenericByteViewArray<T>, ArrowError> {
896    let new_views = take_native(array.views(), indices);
897    let new_nulls = take_nulls::<_, CHECKED>(array.nulls(), indices);
898    let buffers = Arc::clone(array.data_buffers());
899    // Safety:  array.views was valid, and take_native copies only valid values, and verifies bounds
900    Ok(unsafe { GenericByteViewArray::new_unchecked(new_views, buffers, new_nulls) })
901}
902
903/// `take` implementation for list arrays
904///
905/// Copies the selected list entries' child slices into a new child array
906/// via `MutableArrayData`, then reconstructs a list array with new offsets
907fn take_list<IndexType, OffsetType, const CHECKED: bool>(
908    values: &GenericListArray<OffsetType::Native>,
909    indices: &PrimitiveArray<IndexType>,
910) -> Result<GenericListArray<OffsetType::Native>, ArrowError>
911where
912    IndexType: ArrowPrimitiveType,
913    OffsetType: ArrowPrimitiveType,
914    OffsetType::Native: OffsetSizeTrait,
915    PrimitiveArray<OffsetType>: From<Vec<OffsetType::Native>>,
916{
917    let src_offsets = values.value_offsets();
918    let child_data = values.values().to_data();
919    let nulls = take_nulls::<_, CHECKED>(values.nulls(), indices);
920
921    let mut dst_offsets = Vec::with_capacity(indices.len() + 1);
922    dst_offsets.push(OffsetType::Native::zero());
923
924    let field = values.value_field().clone();
925
926    if child_data.null_count() == 0
927        && let Some(bytes_per_value) = child_data.data_type().primitive_width()
928    {
929        let values_buf = &child_data.buffers()[0];
930        let child_buf_offset = child_data.offset() * bytes_per_value;
931
932        let avg_row_len = child_data
933            .len()
934            .checked_div(values.len().max(1))
935            .unwrap_or(0);
936        let mut dst_buf = MutableBuffer::new(
937            avg_row_len
938                .saturating_mul(indices.len())
939                .saturating_mul(bytes_per_value),
940        );
941
942        let mut child_len = OffsetType::Native::zero();
943
944        match nulls.as_ref().filter(|n| n.null_count() > 0) {
945            None => {
946                for &idx in indices.values() {
947                    let row = idx.as_usize();
948                    let start = child_buf_offset + src_offsets[row].as_usize() * bytes_per_value;
949                    let end = child_buf_offset + src_offsets[row + 1].as_usize() * bytes_per_value;
950                    dst_buf.extend_from_slice(&values_buf[start..end]);
951                    child_len = child_len
952                        .checked_add(&(src_offsets[row + 1] - src_offsets[row]))
953                        .ok_or_else(|| ArrowError::OffsetOverflowError(child_len.as_usize()))?;
954                    dst_offsets.push(child_len);
955                }
956            }
957            Some(valid) => {
958                let mut prev = 0;
959                for vidx in valid.valid_indices() {
960                    // Fill offsets for null values between the two valid indices.
961                    if prev < vidx {
962                        dst_offsets.extend(std::iter::repeat_n(child_len, vidx - prev));
963                    }
964                    let row = if CHECKED {
965                        indices.value(vidx).as_usize()
966                    } else {
967                        // SAFETY: !CHECKED means the caller guarantees all indices are valid;
968                        // `vidx` is further bounded by the validity bitmap of `indices`.
969                        unsafe { indices.value_unchecked(vidx) }.as_usize()
970                    };
971                    let start = child_buf_offset + src_offsets[row].as_usize() * bytes_per_value;
972                    let end = child_buf_offset + src_offsets[row + 1].as_usize() * bytes_per_value;
973                    dst_buf.extend_from_slice(&values_buf[start..end]);
974                    child_len = child_len
975                        .checked_add(&(src_offsets[row + 1] - src_offsets[row]))
976                        .ok_or_else(|| ArrowError::OffsetOverflowError(child_len.as_usize()))?;
977                    dst_offsets.push(child_len);
978                    prev = vidx + 1;
979                }
980                dst_offsets.extend(std::iter::repeat_n(child_len, indices.len() - prev));
981            }
982        }
983
984        debug_assert_eq!(
985            dst_offsets.len(),
986            indices.len() + 1,
987            "New offsets was filled under/over the expected capacity"
988        );
989
990        // Safety: data_type, len, and buffer are all derived from the already-validated
991        // source child_data, so re-validation is unnecessary.
992        let child = make_array(unsafe {
993            ArrayData::builder(child_data.data_type().clone())
994                .len(child_len.as_usize())
995                .add_buffer(dst_buf.into())
996                .build_unchecked()
997        });
998        // SAFETY: `dst_offsets` is constructed to be monotonically increasing above.
999        let offsets = unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(dst_offsets)) };
1000        return GenericListArray::<OffsetType::Native>::try_new(field, offsets, child, nulls);
1001    }
1002
1003    let capacity = child_data
1004        .len()
1005        .checked_div(values.len())
1006        .map(|avg| avg * indices.len())
1007        .unwrap_or_default();
1008    let mut mutable =
1009        MutableArrayData::new(vec![&child_data], child_data.null_count() > 0, capacity);
1010
1011    match nulls.as_ref().filter(|n| n.null_count() > 0) {
1012        None => {
1013            for idx in indices.values() {
1014                let row = idx.as_usize();
1015                mutable.try_extend(
1016                    0,
1017                    src_offsets[row].as_usize(),
1018                    src_offsets[row + 1].as_usize(),
1019                )?;
1020                dst_offsets.push(
1021                    OffsetType::Native::from_usize(mutable.len())
1022                        .ok_or_else(|| ArrowError::OffsetOverflowError(mutable.len()))?,
1023                );
1024            }
1025        }
1026        Some(valid) => {
1027            let mut last = 0;
1028            for i in valid.valid_indices() {
1029                let current = OffsetType::Native::from_usize(mutable.len())
1030                    .ok_or_else(|| ArrowError::OffsetOverflowError(mutable.len()))?;
1031                if last < i {
1032                    dst_offsets.extend(std::iter::repeat_n(current, i - last));
1033                }
1034                let row = if CHECKED {
1035                    indices.value(i).as_usize()
1036                } else {
1037                    // SAFETY: !CHECKED means the caller guarantees all indices are valid;
1038                    // `i` is further bounded by the validity bitmap of `indices`.
1039                    unsafe { indices.value_unchecked(i) }.as_usize()
1040                };
1041                mutable.try_extend(
1042                    0,
1043                    src_offsets[row].as_usize(),
1044                    src_offsets[row + 1].as_usize(),
1045                )?;
1046                dst_offsets.push(
1047                    OffsetType::Native::from_usize(mutable.len())
1048                        .ok_or_else(|| ArrowError::OffsetOverflowError(mutable.len()))?,
1049                );
1050                last = i + 1;
1051            }
1052            // Filling offsets for null values at the end
1053            let final_offset = OffsetType::Native::from_usize(mutable.len())
1054                .ok_or_else(|| ArrowError::OffsetOverflowError(mutable.len()))?;
1055            dst_offsets.extend(std::iter::repeat_n(final_offset, indices.len() - last));
1056        }
1057    }
1058
1059    debug_assert_eq!(dst_offsets.len(), indices.len() + 1);
1060
1061    // SAFETY: `dst_offsets` is constructed to be monotonically increasing above
1062    let offsets = unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(dst_offsets)) };
1063    let child = make_array(mutable.freeze());
1064    GenericListArray::<OffsetType::Native>::try_new(field, offsets, child, nulls)
1065}
1066
1067fn take_list_view<IndexType, OffsetType, const CHECKED: bool>(
1068    values: &GenericListViewArray<OffsetType::Native>,
1069    indices: &PrimitiveArray<IndexType>,
1070) -> Result<GenericListViewArray<OffsetType::Native>, ArrowError>
1071where
1072    IndexType: ArrowPrimitiveType,
1073    OffsetType: ArrowPrimitiveType,
1074    OffsetType::Native: OffsetSizeTrait,
1075{
1076    let taken_offsets = take_native(values.offsets(), indices);
1077    let taken_sizes = take_native(values.sizes(), indices);
1078    let nulls = take_nulls::<_, CHECKED>(values.nulls(), indices);
1079
1080    let field = match values.data_type() {
1081        DataType::ListView(field) | DataType::LargeListView(field) => field.clone(),
1082        d => unreachable!("take_list_view called with non-list-view data type {d}"),
1083    };
1084
1085    // SAFETY: the taken offsets/sizes are a permutation of the (valid) input
1086    // offsets/sizes, so they remain within the bounds of the child array.
1087    Ok(unsafe {
1088        GenericListViewArray::<OffsetType::Native>::new_unchecked(
1089            field,
1090            taken_offsets,
1091            taken_sizes,
1092            Arc::clone(values.values()),
1093            nulls,
1094        )
1095    })
1096}
1097
1098/// `take` implementation for `FixedSizeListArray`
1099///
1100/// Calculates the index and indexed offset for the inner array,
1101/// applying `take` on the inner array, then reconstructing a list array
1102/// with the indexed offsets
1103fn take_fixed_size_list<IndexType: ArrowPrimitiveType, const CHECKED: bool>(
1104    values: &FixedSizeListArray,
1105    indices: &PrimitiveArray<IndexType>,
1106    length: <UInt32Type as ArrowPrimitiveType>::Native,
1107) -> Result<FixedSizeListArray, ArrowError> {
1108    let field = values.value_field();
1109    let child = values.values();
1110    let nulls = take_nulls::<_, CHECKED>(values.nulls(), indices);
1111
1112    // Fast path: primitive child with no nulls  copy row-sized byte blocks directly,
1113    let taken_child = if child.null_count() == 0
1114        && let Some(element_size) = child.data_type().primitive_width()
1115    {
1116        take_fixed_size_list_primitive(
1117            child,
1118            indices,
1119            length as usize,
1120            element_size,
1121            nulls.as_ref(),
1122        )
1123    } else {
1124        let list_indices = take_value_indices_from_fixed_size_list(values, indices, length)?;
1125        take_impl::<UInt32Type, CHECKED>(child.as_ref(), &list_indices)?
1126    };
1127
1128    FixedSizeListArray::try_new_with_length(
1129        field.clone(),
1130        length as i32,
1131        taken_child,
1132        nulls,
1133        indices.len(),
1134    )
1135}
1136
1137#[inline(never)]
1138fn take_fixed_size_list_primitive<IndexType: ArrowPrimitiveType>(
1139    child: &ArrayRef,
1140    indices: &PrimitiveArray<IndexType>,
1141    list_size: usize,
1142    element_size: usize,
1143    taken_row_nulls: Option<&NullBuffer>,
1144) -> ArrayRef {
1145    let row_bytes = list_size * element_size;
1146    let child_data = child.to_data();
1147    let src = child_data.buffers()[0].as_slice();
1148    let child_byte_offset = child_data.offset() * element_size;
1149
1150    debug_assert!(
1151        indices.len().checked_mul(row_bytes).is_some(),
1152        "take_fixed_size_list_primitive: output buffer size overflows usize"
1153    );
1154    let out_len = indices.len() * list_size;
1155
1156    let mut out = MutableBuffer::from_len_zeroed(indices.len() * row_bytes);
1157    let out_slice = out.as_slice_mut();
1158
1159    if indices.null_count() == 0 {
1160        for (out_row, index) in indices.values().iter().enumerate() {
1161            let src_start = child_byte_offset + index.as_usize() * row_bytes;
1162            out_slice[out_row * row_bytes..(out_row + 1) * row_bytes]
1163                .copy_from_slice(&src[src_start..src_start + row_bytes]);
1164        }
1165    } else {
1166        for (out_row, index) in indices.values().iter().enumerate() {
1167            if indices.is_valid(out_row) {
1168                let src_start = child_byte_offset + index.as_usize() * row_bytes;
1169                out_slice[out_row * row_bytes..(out_row + 1) * row_bytes]
1170                    .copy_from_slice(&src[src_start..src_start + row_bytes]);
1171            }
1172        }
1173    }
1174    // Expand the per-row null bitmap to per-element: each null row produces `list_size` null elements.
1175    let child_null_buf = taken_row_nulls.map(|n| n.expand(list_size).buffer().clone());
1176
1177    // SAFETY:
1178    // - The data buffer has `indices.len() * row_bytes` bytes = `out_len` elements of a
1179    //   primitive type, matching `.len(out_len)`.
1180    // - The null buffer (when present) is `ceil(out_len, 8)` bytes via `expand(list_size).buffer()`.
1181    // - Primitive types have exactly one value buffer.
1182    let array_data = unsafe {
1183        ArrayData::builder(child.data_type().clone())
1184            .len(out_len)
1185            .add_buffer(out.into())
1186            .null_bit_buffer(child_null_buf)
1187            .build_unchecked()
1188    };
1189    make_array(array_data)
1190}
1191
1192/// The take kernel implementation for `FixedSizeBinaryArray`.
1193///
1194/// The computation is done in two steps:
1195/// - Compute the values buffer
1196/// - Compute the null buffer
1197fn take_fixed_size_binary<IndexType: ArrowPrimitiveType, const CHECKED: bool>(
1198    values: &FixedSizeBinaryArray,
1199    indices: &PrimitiveArray<IndexType>,
1200    size: i32,
1201) -> Result<FixedSizeBinaryArray, ArrowError> {
1202    let size_usize = usize::try_from(size).map_err(|_| {
1203        ArrowError::InvalidArgumentError(format!("Cannot convert size '{size}' to usize"))
1204    })?;
1205
1206    let result_buffer = match size_usize {
1207        1 => take_fixed_size::<IndexType, 1>(values.values(), indices),
1208        2 => take_fixed_size::<IndexType, 2>(values.values(), indices),
1209        4 => take_fixed_size::<IndexType, 4>(values.values(), indices),
1210        8 => take_fixed_size::<IndexType, 8>(values.values(), indices),
1211        16 => take_fixed_size::<IndexType, 16>(values.values(), indices),
1212        _ => take_fixed_size_binary_buffer_dynamic_length(values, indices, size_usize),
1213    };
1214
1215    let value_nulls = take_nulls::<_, CHECKED>(values.nulls(), indices);
1216    let final_nulls = NullBuffer::union(value_nulls.as_ref(), indices.nulls());
1217
1218    return FixedSizeBinaryArray::try_new(size, result_buffer, final_nulls);
1219
1220    /// Implementation of the take kernel for fixed size binary arrays.
1221    #[inline(never)]
1222    fn take_fixed_size_binary_buffer_dynamic_length<IndexType: ArrowPrimitiveType>(
1223        values: &FixedSizeBinaryArray,
1224        indices: &PrimitiveArray<IndexType>,
1225        size_usize: usize,
1226    ) -> Buffer {
1227        let values_buffer = values.values().as_slice();
1228        let mut output = Vec::with_capacity(indices.len() * size_usize);
1229
1230        if indices.null_count() == 0 {
1231            let array_iter = indices.values().iter().map(|idx| {
1232                let offset = idx.as_usize() * size_usize;
1233                &values_buffer[offset..offset + size_usize]
1234            });
1235            for slice in array_iter {
1236                output.extend_from_slice(slice);
1237            }
1238        } else {
1239            // The indices nullability cannot be ignored here because the values buffer may contain
1240            // nulls which should not cause a panic.
1241            let array_iter = indices.iter().map(|idx| {
1242                idx.map(|idx| {
1243                    let offset = idx.as_usize() * size_usize;
1244                    &values_buffer[offset..offset + size_usize]
1245                })
1246            });
1247            for slice in array_iter {
1248                match slice {
1249                    None => output.resize(output.len() + size_usize, 0),
1250                    Some(slice) => output.extend_from_slice(slice),
1251                }
1252            }
1253        }
1254
1255        output.into()
1256    }
1257}
1258
1259/// Implements the take kernel semantics over a flat [`Buffer`], interpreting it as a slice of
1260/// `&[[u8; N]]`, where `N` is a compile-time constant. The usage of a flat [`Buffer`] allows using
1261/// this kernel without an available [`ArrowPrimitiveType`] (e.g., for `[u8; 5]`).
1262///
1263/// # Using This Function in the Primitive Take Kernel
1264///
1265/// This function is basically the same as [`take_native`] but just on a flat [`Buffer`] instead of
1266/// the primitive [`ScalarBuffer`]. Ideally, the [`take_primitive`] kernel should just use this
1267/// more general function. However, the "idiomatic code" requires the
1268/// [feature(generic_const_exprs)](https://github.com/rust-lang/rust/issues/76560) for calling
1269/// `take_fixed_size<I, { size_of::<T::Native> () } >(...)`. Once this feature has been stabilized,
1270/// we can use this function also in the primitive kernels.
1271fn take_fixed_size<IndexType: ArrowPrimitiveType, const N: usize>(
1272    buffer: &Buffer,
1273    indices: &PrimitiveArray<IndexType>,
1274) -> Buffer {
1275    assert_eq!(
1276        buffer.len() % N,
1277        0,
1278        "Invalid array length in take_fixed_size"
1279    );
1280
1281    let ptr = buffer.as_ptr();
1282    let chunk_ptr = ptr.cast::<[u8; N]>();
1283    let chunk_len = buffer.len() / N;
1284    let buffer: &[[u8; N]] = unsafe {
1285        // SAFETY: interpret an already valid slice as a slice of N-byte chunks. N divides buffer
1286        // length without remainder.
1287        std::slice::from_raw_parts(chunk_ptr, chunk_len)
1288    };
1289
1290    let result_buffer = match indices.nulls().filter(|n| n.null_count() > 0) {
1291        Some(n) => indices
1292            .values()
1293            .iter()
1294            .enumerate()
1295            .map(|(idx, index)| match buffer.get(index.as_usize()) {
1296                Some(v) => *v,
1297                // SAFETY: idx<indices.len()
1298                None => match unsafe { n.inner().value_unchecked(idx) } {
1299                    false => [0u8; N],
1300                    true => panic!("Out-of-bounds index {index:?}"),
1301                },
1302            })
1303            .collect::<Vec<_>>(),
1304        None => indices
1305            .values()
1306            .iter()
1307            .map(|index| buffer[index.as_usize()])
1308            .collect::<Vec<_>>(),
1309    };
1310
1311    let mut vec = ManuallyDrop::new(result_buffer); // Prevent de-allocation
1312    let ptr = vec.as_mut_ptr();
1313    let len = vec.len();
1314    let cap = vec.capacity();
1315    let result_buffer = unsafe {
1316        // SAFETY: flattening an already valid Vec.
1317        Vec::from_raw_parts(ptr.cast::<u8>(), len * N, cap * N)
1318    };
1319
1320    Buffer::from_vec(result_buffer)
1321}
1322
1323/// `take` implementation for dictionary arrays
1324///
1325/// applies `take` to the keys of the dictionary array and returns a new dictionary array
1326/// with the same dictionary values and reordered keys
1327fn take_dict<T: ArrowDictionaryKeyType, I: ArrowPrimitiveType, const CHECKED: bool>(
1328    values: &DictionaryArray<T>,
1329    indices: &PrimitiveArray<I>,
1330) -> Result<DictionaryArray<T>, ArrowError> {
1331    let new_keys = take_primitive::<_, _, CHECKED>(values.keys(), indices)?;
1332    Ok(unsafe { DictionaryArray::new_unchecked(new_keys, values.values().clone()) })
1333}
1334
1335/// `take` implementation for run arrays
1336///
1337/// Finds physical indices for the given logical indices and builds output run array
1338/// by taking values in the input run_array.values at the physical indices.
1339/// The output run array will be run encoded on the physical indices and not on output values.
1340/// For e.g. an input `RunArray{ run_ends = [2,4,6,8], values=[1,2,1,2] }` and `logical_indices=[2,3,6,7]`
1341/// would be converted to `physical_indices=[1,1,3,3]` which will be used to build
1342/// output `RunArray{ run_ends=[2,4], values=[2,2] }`.
1343///
1344/// A null logical index becomes a null run. Consecutive nulls are merged.
1345fn take_run<T: RunEndIndexType, I: ArrowPrimitiveType>(
1346    run_array: &RunArray<T>,
1347    logical_indices: &PrimitiveArray<I>,
1348) -> Result<RunArray<T>, ArrowError> {
1349    let physical_indices = physical_indices_for_take(run_array, logical_indices)?;
1350
1351    // Run encode the physical indices into new_run_ends
1352    // Keep track of the physical indices to take in take_value_indices
1353    // `unwrap` is used in this function because the unwrapped values are bounded by the corresponding `::Native`.
1354    let mut new_run_ends = Vec::with_capacity(1);
1355    let mut take_value_indices = Vec::with_capacity(1);
1356    let mut take_value_is_valid = NullBufferBuilder::new(1);
1357
1358    let values_cmp = make_comparator(
1359        run_array.values().as_ref(),
1360        run_array.values().as_ref(),
1361        SortOptions::default(),
1362    )?;
1363
1364    for ix in 1..physical_indices.len() {
1365        let prev_idx = physical_indices[ix - 1];
1366        let cur_idx = physical_indices[ix];
1367        if is_new_run_take(run_array.values().as_ref(), prev_idx, cur_idx, &values_cmp) {
1368            // Safe unwrap since physical indices came from a valid run array.
1369            let index = I::Native::from_usize(prev_idx.unwrap_or_default()).unwrap();
1370            take_value_indices.push(index);
1371            take_value_is_valid
1372                .append(prev_idx.is_some_and(|idx| run_array.values().is_valid(idx)));
1373            new_run_ends.push(T::Native::from_usize(ix).unwrap());
1374        }
1375    }
1376    let last = physical_indices[physical_indices.len() - 1];
1377    // Safe unwrap since physical indices came from a valid run array.
1378    let index = I::Native::from_usize(last.unwrap_or_default()).unwrap();
1379    take_value_indices.push(index);
1380    take_value_is_valid.append(last.is_some_and(|idx| run_array.values().is_valid(idx)));
1381    new_run_ends.push(T::Native::from_usize(physical_indices.len()).unwrap());
1382
1383    // SAFETY: run-ends are strictly increasing with last value == logical length.
1384    let run_ends = unsafe {
1385        RunEndBuffer::new_unchecked(ScalarBuffer::from(new_run_ends), 0, physical_indices.len())
1386    };
1387
1388    let nulls = take_value_is_valid.finish();
1389    let take_value_indices =
1390        PrimitiveArray::<I>::new(ScalarBuffer::from(take_value_indices), nulls);
1391
1392    let new_values = take(run_array.values(), &take_value_indices, None)?;
1393
1394    // SAFETY: `new_values` has one entry per run.
1395    Ok(
1396        unsafe {
1397            RunArray::<T>::new_unchecked(run_array.data_type().clone(), run_ends, new_values)
1398        },
1399    )
1400}
1401
1402/// Physical run index for each logical take slot.
1403///
1404/// `None` means the logical index is null. Only valid indices are passed to
1405/// [`RunArray::get_physical_indices`]; a null slot's backing integer is ignored
1406/// and may be out of range.
1407fn physical_indices_for_take<T: RunEndIndexType, I: ArrowPrimitiveType>(
1408    run_array: &RunArray<T>,
1409    logical_indices: &PrimitiveArray<I>,
1410) -> Result<Vec<Option<usize>>, ArrowError> {
1411    if logical_indices.null_count() == 0 {
1412        return Ok(run_array
1413            .get_physical_indices(logical_indices.values())?
1414            .into_iter()
1415            .map(Some)
1416            .collect());
1417    }
1418
1419    let valid_logical: Vec<_> = logical_indices.iter().flatten().collect();
1420
1421    let valid_physical = if valid_logical.is_empty() {
1422        Vec::new()
1423    } else {
1424        run_array.get_physical_indices(&valid_logical)?
1425    };
1426
1427    let mut valid_physical = valid_physical.into_iter();
1428    Ok(logical_indices
1429        .iter()
1430        .map(|index| index.map(|_| valid_physical.next().unwrap()))
1431        .collect())
1432}
1433
1434fn is_new_run_take(
1435    values: &dyn Array,
1436    prev_idx: Option<usize>,
1437    cur_idx: Option<usize>,
1438    values_cmp: &arrow_cmp::DynComparator,
1439) -> bool {
1440    let prev_valid = prev_idx.is_some_and(|idx| values.is_valid(idx));
1441    let cur_valid = cur_idx.is_some_and(|idx| values.is_valid(idx));
1442    match (prev_valid, cur_valid) {
1443        (false, false) => false,
1444        (true, true) => {
1445            let prev = prev_idx.unwrap();
1446            let cur = cur_idx.unwrap();
1447            prev != cur && values_cmp(cur, prev).is_ne()
1448        }
1449        _ => true,
1450    }
1451}
1452
1453/// Takes/filters a fixed size list array's inner data using the offsets of the list array.
1454fn take_value_indices_from_fixed_size_list<IndexType>(
1455    list: &FixedSizeListArray,
1456    indices: &PrimitiveArray<IndexType>,
1457    length: <UInt32Type as ArrowPrimitiveType>::Native,
1458) -> Result<PrimitiveArray<UInt32Type>, ArrowError>
1459where
1460    IndexType: ArrowPrimitiveType,
1461{
1462    let mut values = UInt32Builder::with_capacity(length as usize * indices.len());
1463
1464    for i in 0..indices.len() {
1465        if indices.is_valid(i) {
1466            let index = indices
1467                .value(i)
1468                .to_usize()
1469                .ok_or_else(|| ArrowError::ComputeError("Cast to usize failed".to_string()))?;
1470            let start = list.value_offset(index) as <UInt32Type as ArrowPrimitiveType>::Native;
1471
1472            // Safety: Range always has known length.
1473            unsafe {
1474                values.append_trusted_len_iter(start..start + length);
1475            }
1476        } else {
1477            values.append_nulls(length as usize);
1478        }
1479    }
1480
1481    Ok(values.finish())
1482}
1483
1484/// To avoid generating take implementations for every index type, instead we
1485/// only generate for UInt32 and UInt64 and coerce inputs to these types
1486trait ToIndices {
1487    type T: ArrowPrimitiveType;
1488
1489    fn to_indices(&self) -> PrimitiveArray<Self::T>;
1490}
1491
1492macro_rules! to_indices_reinterpret {
1493    ($t:ty, $o:ty) => {
1494        impl ToIndices for PrimitiveArray<$t> {
1495            type T = $o;
1496
1497            fn to_indices(&self) -> PrimitiveArray<$o> {
1498                let cast = ScalarBuffer::new(self.values().inner().clone(), 0, self.len());
1499                PrimitiveArray::new(cast, self.nulls().cloned())
1500            }
1501        }
1502    };
1503}
1504
1505macro_rules! to_indices_identity {
1506    ($t:ty) => {
1507        impl ToIndices for PrimitiveArray<$t> {
1508            type T = $t;
1509
1510            fn to_indices(&self) -> PrimitiveArray<$t> {
1511                self.clone()
1512            }
1513        }
1514    };
1515}
1516
1517macro_rules! to_indices_widening {
1518    ($t:ty, $o:ty) => {
1519        impl ToIndices for PrimitiveArray<$t> {
1520            type T = UInt32Type;
1521
1522            fn to_indices(&self) -> PrimitiveArray<$o> {
1523                let cast = self.values().iter().copied().map(|x| x as _).collect();
1524                PrimitiveArray::new(cast, self.nulls().cloned())
1525            }
1526        }
1527    };
1528}
1529
1530to_indices_widening!(UInt8Type, UInt32Type);
1531to_indices_widening!(Int8Type, UInt32Type);
1532
1533to_indices_widening!(UInt16Type, UInt32Type);
1534to_indices_widening!(Int16Type, UInt32Type);
1535
1536to_indices_identity!(UInt32Type);
1537to_indices_reinterpret!(Int32Type, UInt32Type);
1538
1539to_indices_identity!(UInt64Type);
1540to_indices_reinterpret!(Int64Type, UInt64Type);
1541
1542/// Take rows by index from [`RecordBatch`] and returns a new [`RecordBatch`] from those indexes.
1543///
1544/// This function will call [`take`] on each array of the [`RecordBatch`] and assemble a new [`RecordBatch`].
1545///
1546/// # Example
1547/// ```
1548/// # use std::sync::Arc;
1549/// # use arrow_array::{StringArray, Int32Array, UInt32Array, RecordBatch};
1550/// # use arrow_schema::{DataType, Field, Schema};
1551/// # use arrow_select::take::take_record_batch;
1552/// let schema = Arc::new(Schema::new(vec![
1553///     Field::new("a", DataType::Int32, true),
1554///     Field::new("b", DataType::Utf8, true),
1555/// ]));
1556/// let batch = RecordBatch::try_new(
1557///     schema.clone(),
1558///     vec![
1559///         Arc::new(Int32Array::from_iter_values(0..20)),
1560///         Arc::new(StringArray::from_iter_values(
1561///             (0..20).map(|i| format!("str-{}", i)),
1562///         )),
1563///     ],
1564/// )
1565/// .unwrap();
1566///
1567/// let indices = UInt32Array::from(vec![1, 5, 10]);
1568/// let taken = take_record_batch(&batch, &indices).unwrap();
1569///
1570/// let expected = RecordBatch::try_new(
1571///     schema,
1572///     vec![
1573///         Arc::new(Int32Array::from(vec![1, 5, 10])),
1574///         Arc::new(StringArray::from(vec!["str-1", "str-5", "str-10"])),
1575///     ],
1576/// )
1577/// .unwrap();
1578/// assert_eq!(taken, expected);
1579/// ```
1580pub fn take_record_batch(
1581    record_batch: &RecordBatch,
1582    indices: &dyn Array,
1583) -> Result<RecordBatch, ArrowError> {
1584    let columns = record_batch
1585        .columns()
1586        .iter()
1587        .map(|c| take(c, indices, None))
1588        .collect::<Result<Vec<_>, _>>()?;
1589    RecordBatch::try_new(record_batch.schema(), columns)
1590}
1591
1592#[cfg(test)]
1593mod tests {
1594    use super::*;
1595    use arrow_array::builder::*;
1596    use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano};
1597    use arrow_data::ArrayData;
1598    use arrow_schema::{Field, Fields, TimeUnit, UnionFields};
1599    use num_traits::ToPrimitive;
1600
1601    fn test_take_decimal_arrays(
1602        data: Vec<Option<i128>>,
1603        index: &UInt32Array,
1604        options: Option<TakeOptions>,
1605        expected_data: Vec<Option<i128>>,
1606        precision: &u8,
1607        scale: &i8,
1608    ) -> Result<(), ArrowError> {
1609        let output = data
1610            .into_iter()
1611            .collect::<Decimal128Array>()
1612            .with_precision_and_scale(*precision, *scale)
1613            .unwrap();
1614
1615        let expected = expected_data
1616            .into_iter()
1617            .collect::<Decimal128Array>()
1618            .with_precision_and_scale(*precision, *scale)
1619            .unwrap();
1620
1621        let expected = Arc::new(expected) as ArrayRef;
1622        let output = take(&output, index, options).unwrap();
1623        assert_eq!(&output, &expected);
1624        Ok(())
1625    }
1626
1627    fn test_take_boolean_arrays(
1628        data: Vec<Option<bool>>,
1629        index: &UInt32Array,
1630        options: Option<TakeOptions>,
1631        expected_data: Vec<Option<bool>>,
1632    ) {
1633        let output = BooleanArray::from(data);
1634        let expected = Arc::new(BooleanArray::from(expected_data)) as ArrayRef;
1635        let output = take(&output, index, options).unwrap();
1636        assert_eq!(&output, &expected)
1637    }
1638
1639    fn test_take_primitive_arrays<T>(
1640        data: Vec<Option<T::Native>>,
1641        index: &UInt32Array,
1642        options: Option<TakeOptions>,
1643        expected_data: Vec<Option<T::Native>>,
1644    ) -> Result<(), ArrowError>
1645    where
1646        T: ArrowPrimitiveType,
1647        PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
1648    {
1649        let output = PrimitiveArray::<T>::from(data);
1650        let expected = Arc::new(PrimitiveArray::<T>::from(expected_data)) as ArrayRef;
1651        let output = take(&output, index, options)?;
1652        assert_eq!(&output, &expected);
1653        Ok(())
1654    }
1655
1656    fn test_take_primitive_arrays_non_null<T>(
1657        data: Vec<T::Native>,
1658        index: &UInt32Array,
1659        options: Option<TakeOptions>,
1660        expected_data: Vec<Option<T::Native>>,
1661    ) -> Result<(), ArrowError>
1662    where
1663        T: ArrowPrimitiveType,
1664        PrimitiveArray<T>: From<Vec<T::Native>>,
1665        PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
1666    {
1667        let output = PrimitiveArray::<T>::from(data);
1668        let expected = Arc::new(PrimitiveArray::<T>::from(expected_data)) as ArrayRef;
1669        let output = take(&output, index, options)?;
1670        assert_eq!(&output, &expected);
1671        Ok(())
1672    }
1673
1674    fn test_take_impl_primitive_arrays<T, I>(
1675        data: Vec<Option<T::Native>>,
1676        index: &PrimitiveArray<I>,
1677        options: Option<TakeOptions>,
1678        expected_data: Vec<Option<T::Native>>,
1679    ) where
1680        T: ArrowPrimitiveType,
1681        PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
1682        I: ArrowPrimitiveType,
1683    {
1684        let output = PrimitiveArray::<T>::from(data);
1685        let expected = PrimitiveArray::<T>::from(expected_data);
1686        let output = take(&output, index, options).unwrap();
1687        let output = output.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();
1688        assert_eq!(output, &expected)
1689    }
1690
1691    // create a simple struct for testing purposes
1692    fn create_test_struct(values: Vec<Option<(Option<bool>, Option<i32>)>>) -> StructArray {
1693        let mut struct_builder = StructBuilder::new(
1694            Fields::from(vec![
1695                Field::new("a", DataType::Boolean, true),
1696                Field::new("b", DataType::Int32, true),
1697            ]),
1698            vec![
1699                Box::new(BooleanBuilder::with_capacity(values.len())),
1700                Box::new(Int32Builder::with_capacity(values.len())),
1701            ],
1702        );
1703
1704        for value in values {
1705            struct_builder
1706                .field_builder::<BooleanBuilder>(0)
1707                .unwrap()
1708                .append_option(value.and_then(|v| v.0));
1709            struct_builder
1710                .field_builder::<Int32Builder>(1)
1711                .unwrap()
1712                .append_option(value.and_then(|v| v.1));
1713            struct_builder.append(value.is_some());
1714        }
1715        struct_builder.finish()
1716    }
1717
1718    #[test]
1719    fn test_take_decimal128_non_null_indices() {
1720        let index = UInt32Array::from(vec![0, 5, 3, 1, 4, 2]);
1721        let precision: u8 = 10;
1722        let scale: i8 = 5;
1723        test_take_decimal_arrays(
1724            vec![None, Some(3), Some(5), Some(2), Some(3), None],
1725            &index,
1726            None,
1727            vec![None, None, Some(2), Some(3), Some(3), Some(5)],
1728            &precision,
1729            &scale,
1730        )
1731        .unwrap();
1732    }
1733
1734    #[test]
1735    fn test_take_decimal128() {
1736        let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
1737        let precision: u8 = 10;
1738        let scale: i8 = 5;
1739        test_take_decimal_arrays(
1740            vec![Some(0), Some(1), Some(2), Some(3), Some(4)],
1741            &index,
1742            None,
1743            vec![Some(3), None, Some(1), Some(3), Some(2)],
1744            &precision,
1745            &scale,
1746        )
1747        .unwrap();
1748    }
1749
1750    #[test]
1751    fn test_take_primitive_non_null_indices() {
1752        let index = UInt32Array::from(vec![0, 5, 3, 1, 4, 2]);
1753        test_take_primitive_arrays::<Int8Type>(
1754            vec![None, Some(3), Some(5), Some(2), Some(3), None],
1755            &index,
1756            None,
1757            vec![None, None, Some(2), Some(3), Some(3), Some(5)],
1758        )
1759        .unwrap();
1760    }
1761
1762    #[test]
1763    fn test_take_primitive_non_null_values() {
1764        let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
1765        test_take_primitive_arrays::<Int8Type>(
1766            vec![Some(0), Some(1), Some(2), Some(3), Some(4)],
1767            &index,
1768            None,
1769            vec![Some(3), None, Some(1), Some(3), Some(2)],
1770        )
1771        .unwrap();
1772    }
1773
1774    #[test]
1775    fn test_take_primitive_non_null() {
1776        let index = UInt32Array::from(vec![0, 5, 3, 1, 4, 2]);
1777        test_take_primitive_arrays::<Int8Type>(
1778            vec![Some(0), Some(3), Some(5), Some(2), Some(3), Some(1)],
1779            &index,
1780            None,
1781            vec![Some(0), Some(1), Some(2), Some(3), Some(3), Some(5)],
1782        )
1783        .unwrap();
1784    }
1785
1786    #[test]
1787    fn test_take_primitive_nullable_indices_non_null_values_with_offset() {
1788        let index = UInt32Array::from(vec![Some(0), Some(1), Some(2), Some(3), None, None]);
1789        let index = index.slice(2, 4);
1790        let index = index.as_any().downcast_ref::<UInt32Array>().unwrap();
1791
1792        assert_eq!(
1793            index,
1794            &UInt32Array::from(vec![Some(2), Some(3), None, None])
1795        );
1796
1797        test_take_primitive_arrays_non_null::<Int64Type>(
1798            vec![0, 10, 20, 30, 40, 50],
1799            index,
1800            None,
1801            vec![Some(20), Some(30), None, None],
1802        )
1803        .unwrap();
1804    }
1805
1806    #[test]
1807    fn test_take_primitive_nullable_indices_nullable_values_with_offset() {
1808        let index = UInt32Array::from(vec![Some(0), Some(1), Some(2), Some(3), None, None]);
1809        let index = index.slice(2, 4);
1810        let index = index.as_any().downcast_ref::<UInt32Array>().unwrap();
1811
1812        assert_eq!(
1813            index,
1814            &UInt32Array::from(vec![Some(2), Some(3), None, None])
1815        );
1816
1817        test_take_primitive_arrays::<Int64Type>(
1818            vec![None, None, Some(20), Some(30), Some(40), Some(50)],
1819            index,
1820            None,
1821            vec![Some(20), Some(30), None, None],
1822        )
1823        .unwrap();
1824    }
1825
1826    #[test]
1827    fn test_take_primitive() {
1828        let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
1829
1830        // int8
1831        test_take_primitive_arrays::<Int8Type>(
1832            vec![Some(0), None, Some(2), Some(3), None],
1833            &index,
1834            None,
1835            vec![Some(3), None, None, Some(3), Some(2)],
1836        )
1837        .unwrap();
1838
1839        // int16
1840        test_take_primitive_arrays::<Int16Type>(
1841            vec![Some(0), None, Some(2), Some(3), None],
1842            &index,
1843            None,
1844            vec![Some(3), None, None, Some(3), Some(2)],
1845        )
1846        .unwrap();
1847
1848        // int32
1849        test_take_primitive_arrays::<Int32Type>(
1850            vec![Some(0), None, Some(2), Some(3), None],
1851            &index,
1852            None,
1853            vec![Some(3), None, None, Some(3), Some(2)],
1854        )
1855        .unwrap();
1856
1857        // int64
1858        test_take_primitive_arrays::<Int64Type>(
1859            vec![Some(0), None, Some(2), Some(3), None],
1860            &index,
1861            None,
1862            vec![Some(3), None, None, Some(3), Some(2)],
1863        )
1864        .unwrap();
1865
1866        // uint8
1867        test_take_primitive_arrays::<UInt8Type>(
1868            vec![Some(0), None, Some(2), Some(3), None],
1869            &index,
1870            None,
1871            vec![Some(3), None, None, Some(3), Some(2)],
1872        )
1873        .unwrap();
1874
1875        // uint16
1876        test_take_primitive_arrays::<UInt16Type>(
1877            vec![Some(0), None, Some(2), Some(3), None],
1878            &index,
1879            None,
1880            vec![Some(3), None, None, Some(3), Some(2)],
1881        )
1882        .unwrap();
1883
1884        // uint32
1885        test_take_primitive_arrays::<UInt32Type>(
1886            vec![Some(0), None, Some(2), Some(3), None],
1887            &index,
1888            None,
1889            vec![Some(3), None, None, Some(3), Some(2)],
1890        )
1891        .unwrap();
1892
1893        // int64
1894        test_take_primitive_arrays::<Int64Type>(
1895            vec![Some(0), None, Some(2), Some(-15), None],
1896            &index,
1897            None,
1898            vec![Some(-15), None, None, Some(-15), Some(2)],
1899        )
1900        .unwrap();
1901
1902        // interval_year_month
1903        test_take_primitive_arrays::<IntervalYearMonthType>(
1904            vec![Some(0), None, Some(2), Some(-15), None],
1905            &index,
1906            None,
1907            vec![Some(-15), None, None, Some(-15), Some(2)],
1908        )
1909        .unwrap();
1910
1911        // interval_day_time
1912        let v1 = IntervalDayTime::new(0, 0);
1913        let v2 = IntervalDayTime::new(2, 0);
1914        let v3 = IntervalDayTime::new(-15, 0);
1915        test_take_primitive_arrays::<IntervalDayTimeType>(
1916            vec![Some(v1), None, Some(v2), Some(v3), None],
1917            &index,
1918            None,
1919            vec![Some(v3), None, None, Some(v3), Some(v2)],
1920        )
1921        .unwrap();
1922
1923        // interval_month_day_nano
1924        let v1 = IntervalMonthDayNano::new(0, 0, 0);
1925        let v2 = IntervalMonthDayNano::new(2, 0, 0);
1926        let v3 = IntervalMonthDayNano::new(-15, 0, 0);
1927        test_take_primitive_arrays::<IntervalMonthDayNanoType>(
1928            vec![Some(v1), None, Some(v2), Some(v3), None],
1929            &index,
1930            None,
1931            vec![Some(v3), None, None, Some(v3), Some(v2)],
1932        )
1933        .unwrap();
1934
1935        // duration_second
1936        test_take_primitive_arrays::<DurationSecondType>(
1937            vec![Some(0), None, Some(2), Some(-15), None],
1938            &index,
1939            None,
1940            vec![Some(-15), None, None, Some(-15), Some(2)],
1941        )
1942        .unwrap();
1943
1944        // duration_millisecond
1945        test_take_primitive_arrays::<DurationMillisecondType>(
1946            vec![Some(0), None, Some(2), Some(-15), None],
1947            &index,
1948            None,
1949            vec![Some(-15), None, None, Some(-15), Some(2)],
1950        )
1951        .unwrap();
1952
1953        // duration_microsecond
1954        test_take_primitive_arrays::<DurationMicrosecondType>(
1955            vec![Some(0), None, Some(2), Some(-15), None],
1956            &index,
1957            None,
1958            vec![Some(-15), None, None, Some(-15), Some(2)],
1959        )
1960        .unwrap();
1961
1962        // duration_nanosecond
1963        test_take_primitive_arrays::<DurationNanosecondType>(
1964            vec![Some(0), None, Some(2), Some(-15), None],
1965            &index,
1966            None,
1967            vec![Some(-15), None, None, Some(-15), Some(2)],
1968        )
1969        .unwrap();
1970
1971        // float32
1972        test_take_primitive_arrays::<Float32Type>(
1973            vec![Some(0.0), None, Some(2.21), Some(-3.1), None],
1974            &index,
1975            None,
1976            vec![Some(-3.1), None, None, Some(-3.1), Some(2.21)],
1977        )
1978        .unwrap();
1979
1980        // float64
1981        test_take_primitive_arrays::<Float64Type>(
1982            vec![Some(0.0), None, Some(2.21), Some(-3.1), None],
1983            &index,
1984            None,
1985            vec![Some(-3.1), None, None, Some(-3.1), Some(2.21)],
1986        )
1987        .unwrap();
1988    }
1989
1990    #[test]
1991    fn test_take_preserve_timezone() {
1992        let index = Int64Array::from(vec![Some(0), None]);
1993
1994        let input = TimestampNanosecondArray::from(vec![
1995            1_639_715_368_000_000_000,
1996            1_639_715_368_000_000_000,
1997        ])
1998        .with_timezone("UTC".to_string());
1999        let result = take(&input, &index, None).unwrap();
2000        match result.data_type() {
2001            DataType::Timestamp(TimeUnit::Nanosecond, tz) => {
2002                assert_eq!(tz.clone(), Some("UTC".into()))
2003            }
2004            _ => panic!(),
2005        }
2006    }
2007
2008    #[test]
2009    fn test_take_impl_primitive_with_int64_indices() {
2010        let index = Int64Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
2011
2012        // int16
2013        test_take_impl_primitive_arrays::<Int16Type, Int64Type>(
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
2020        // int64
2021        test_take_impl_primitive_arrays::<Int64Type, Int64Type>(
2022            vec![Some(0), None, Some(2), Some(-15), None],
2023            &index,
2024            None,
2025            vec![Some(-15), None, None, Some(-15), Some(2)],
2026        );
2027
2028        // uint64
2029        test_take_impl_primitive_arrays::<UInt64Type, Int64Type>(
2030            vec![Some(0), None, Some(2), Some(3), None],
2031            &index,
2032            None,
2033            vec![Some(3), None, None, Some(3), Some(2)],
2034        );
2035
2036        // duration_millisecond
2037        test_take_impl_primitive_arrays::<DurationMillisecondType, Int64Type>(
2038            vec![Some(0), None, Some(2), Some(-15), None],
2039            &index,
2040            None,
2041            vec![Some(-15), None, None, Some(-15), Some(2)],
2042        );
2043
2044        // float32
2045        test_take_impl_primitive_arrays::<Float32Type, Int64Type>(
2046            vec![Some(0.0), None, Some(2.21), Some(-3.1), None],
2047            &index,
2048            None,
2049            vec![Some(-3.1), None, None, Some(-3.1), Some(2.21)],
2050        );
2051    }
2052
2053    #[test]
2054    fn test_take_impl_primitive_with_uint8_indices() {
2055        let index = UInt8Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
2056
2057        // int16
2058        test_take_impl_primitive_arrays::<Int16Type, UInt8Type>(
2059            vec![Some(0), None, Some(2), Some(3), None],
2060            &index,
2061            None,
2062            vec![Some(3), None, None, Some(3), Some(2)],
2063        );
2064
2065        // duration_millisecond
2066        test_take_impl_primitive_arrays::<DurationMillisecondType, UInt8Type>(
2067            vec![Some(0), None, Some(2), Some(-15), None],
2068            &index,
2069            None,
2070            vec![Some(-15), None, None, Some(-15), Some(2)],
2071        );
2072
2073        // float32
2074        test_take_impl_primitive_arrays::<Float32Type, UInt8Type>(
2075            vec![Some(0.0), None, Some(2.21), Some(-3.1), None],
2076            &index,
2077            None,
2078            vec![Some(-3.1), None, None, Some(-3.1), Some(2.21)],
2079        );
2080    }
2081
2082    #[test]
2083    fn test_take_bool() {
2084        let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);
2085        // boolean
2086        test_take_boolean_arrays(
2087            vec![Some(false), None, Some(true), Some(false), None],
2088            &index,
2089            None,
2090            vec![Some(false), None, None, Some(false), Some(true)],
2091        );
2092    }
2093
2094    #[test]
2095    fn test_take_bool_nullable_index() {
2096        // indices where the masked invalid elements would be out of bounds
2097        let index_data = ArrayData::try_new(
2098            DataType::UInt32,
2099            6,
2100            Some(Buffer::from_iter(vec![
2101                false, true, false, true, false, true,
2102            ])),
2103            0,
2104            vec![Buffer::from_iter(vec![99, 0, 999, 1, 9999, 2])],
2105            vec![],
2106        )
2107        .unwrap();
2108        let index = UInt32Array::from(index_data);
2109        test_take_boolean_arrays(
2110            vec![Some(true), None, Some(false)],
2111            &index,
2112            None,
2113            vec![None, Some(true), None, None, None, Some(false)],
2114        );
2115    }
2116
2117    #[test]
2118    fn test_take_bool_nullable_index_nonnull_values() {
2119        // indices where the masked invalid elements would be out of bounds
2120        let index_data = ArrayData::try_new(
2121            DataType::UInt32,
2122            6,
2123            Some(Buffer::from_iter(vec![
2124                false, true, false, true, false, true,
2125            ])),
2126            0,
2127            vec![Buffer::from_iter(vec![99, 0, 999, 1, 9999, 2])],
2128            vec![],
2129        )
2130        .unwrap();
2131        let index = UInt32Array::from(index_data);
2132        test_take_boolean_arrays(
2133            vec![Some(true), Some(true), Some(false)],
2134            &index,
2135            None,
2136            vec![None, Some(true), None, Some(true), None, Some(false)],
2137        );
2138    }
2139
2140    #[test]
2141    fn test_take_bool_with_offset() {
2142        let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(2), None]);
2143        let index = index.slice(2, 4);
2144        let index = index
2145            .as_any()
2146            .downcast_ref::<PrimitiveArray<UInt32Type>>()
2147            .unwrap();
2148
2149        // boolean
2150        test_take_boolean_arrays(
2151            vec![Some(false), None, Some(true), Some(false), None],
2152            index,
2153            None,
2154            vec![None, Some(false), Some(true), None],
2155        );
2156    }
2157
2158    #[test]
2159    // Null indices + sliced source boolean array — exercises src_offset in the sparse take_bits path.
2160    fn test_take_bool_nullable_index_sliced_source() {
2161        let source = BooleanArray::from(vec![Some(true), Some(false), Some(true), Some(false)]);
2162        let source = source.slice(1, 3); // logical: [false, true, false], offset=1
2163        let source = source;
2164
2165        let indices = UInt32Array::from(vec![Some(2), None, Some(0)]);
2166        let result = take(&source, &indices, None).unwrap();
2167        let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
2168
2169        let expected = BooleanArray::from(vec![Some(false), None, Some(false)]);
2170        assert_eq!(result, &expected);
2171    }
2172
2173    #[test]
2174    // >8 elements: exercises the 8-at-a-time unrolled byte packing in take_bits.
2175    fn test_take_bool_no_nulls_multi_byte() {
2176        let source = BooleanArray::from(vec![
2177            true, false, true, true, false, false, true, false, true, true,
2178        ]);
2179        let indices = UInt32Array::from(vec![0, 2, 4, 6, 8, 1, 3, 5, 7, 9]);
2180        let result = take(&source, &indices, None).unwrap();
2181        let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
2182        let expected = BooleanArray::from(vec![
2183            true, true, false, true, true, false, true, false, false, true,
2184        ]);
2185        assert_eq!(result, &expected);
2186    }
2187
2188    #[test]
2189    // Sliced source: verifies src_offset is applied when no null indices.
2190    fn test_take_bool_no_nulls_sliced_source() {
2191        let source = BooleanArray::from(vec![true, false, true, false, true]);
2192        let source = source.slice(2, 3); // [true, false, true], offset=2
2193        let source = source.as_any().downcast_ref::<BooleanArray>().unwrap();
2194        let indices = UInt32Array::from(vec![2, 0, 1]);
2195        let result = take(source, &indices, None).unwrap();
2196        let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
2197        let expected = BooleanArray::from(vec![true, true, false]);
2198        assert_eq!(result, &expected);
2199    }
2200
2201    #[test]
2202    // Nullable source, >8 elements: exercises the 8-at-a-time unrolled byte packing in take_bits_with_validity.
2203    fn test_take_bool_nullable_values_multi_byte() {
2204        let source = BooleanArray::from(vec![
2205            Some(true),
2206            None,
2207            Some(false),
2208            Some(true),
2209            None,
2210            Some(false),
2211            Some(true),
2212            Some(false),
2213            Some(true),
2214            None,
2215        ]);
2216        let indices = UInt32Array::from(vec![0, 2, 4, 6, 8, 1, 3, 5, 7, 9]);
2217        let result = take(&source, &indices, None).unwrap();
2218        let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
2219        let expected = BooleanArray::from(vec![
2220            Some(true),
2221            Some(false),
2222            None,
2223            Some(true),
2224            Some(true),
2225            None,
2226            Some(true),
2227            Some(false),
2228            Some(false),
2229            None,
2230        ]);
2231        assert_eq!(result, &expected);
2232    }
2233
2234    #[test]
2235    // Nullable source, null indices, sliced source: verifies src_offset with both null paths.
2236    fn test_take_bool_nullable_values_sliced_source_null_indices() {
2237        let source =
2238            BooleanArray::from(vec![Some(true), Some(false), None, Some(true), Some(false)]);
2239        let source = source.slice(1, 4); // [false, null, true, false], offset=1
2240        let source = source.as_any().downcast_ref::<BooleanArray>().unwrap();
2241        let indices = UInt32Array::from(vec![Some(3), None, Some(1), Some(0)]);
2242        let result = take(source, &indices, None).unwrap();
2243        let result = result.as_any().downcast_ref::<BooleanArray>().unwrap();
2244        let expected = BooleanArray::from(vec![Some(false), None, None, Some(false)]);
2245        assert_eq!(result, &expected);
2246    }
2247
2248    #[test]
2249    #[should_panic(expected = "assertion failed: idx < self.bit_len")]
2250    fn test_take_bool_oob_no_check_bounds_panics() {
2251        let array = BooleanArray::from(vec![true, false, true]);
2252        let indices = Int32Array::from(vec![0, 1, 10]);
2253        take(&array, &indices, None).unwrap();
2254    }
2255
2256    fn _test_take_string<'a, K>()
2257    where
2258        K: Array + PartialEq + From<Vec<Option<&'a str>>> + 'static,
2259    {
2260        let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(4)]);
2261
2262        let array = K::from(vec![
2263            Some("one"),
2264            None,
2265            Some("three"),
2266            Some("four"),
2267            Some("five"),
2268        ]);
2269        let actual = take(&array, &index, None).unwrap();
2270        assert_eq!(actual.len(), index.len());
2271
2272        let actual = actual.as_any().downcast_ref::<K>().unwrap();
2273
2274        let expected = K::from(vec![Some("four"), None, None, Some("four"), Some("five")]);
2275
2276        assert_eq!(actual, &expected);
2277    }
2278
2279    #[test]
2280    fn test_take_string() {
2281        _test_take_string::<StringArray>()
2282    }
2283
2284    #[test]
2285    fn test_take_large_string() {
2286        _test_take_string::<LargeStringArray>()
2287    }
2288
2289    #[test]
2290    fn test_take_slice_string() {
2291        let strings = StringArray::from(vec![Some("hello"), None, Some("world"), None, Some("hi")]);
2292        let indices = Int32Array::from(vec![Some(0), Some(1), None, Some(0), Some(2)]);
2293        let indices_slice = indices.slice(1, 4);
2294        let expected = StringArray::from(vec![None, None, Some("hello"), Some("world")]);
2295        let result = take(&strings, &indices_slice, None).unwrap();
2296        assert_eq!(result.as_ref(), &expected);
2297    }
2298
2299    /// Take from a *sliced* byte array, i.e. one whose value offsets do not
2300    /// start at zero. This exercises copying byte data out of an array with a
2301    /// non-zero base offset for both the no-null fast path and the nullable
2302    /// path (null indices and selected null values).
2303    #[test]
2304    fn test_take_bytes_sliced_values() {
2305        let values = StringArray::from(vec![
2306            Some("aaa"),
2307            Some("bbb"),
2308            None,
2309            Some("ccccc"),
2310            Some("dd"),
2311            None,
2312            Some("eeee"),
2313        ]);
2314        // Slice so the underlying value offsets no longer start at 0:
2315        // sliced == [None, "ccccc", "dd", None, "eeee"]
2316        let sliced = values.slice(2, 5);
2317
2318        // Fast path: every output slot is valid (no null indices, no null
2319        // values selected).
2320        let indices = Int32Array::from(vec![1, 2, 4, 1]);
2321        let result = take(&sliced, &indices, None).unwrap();
2322        let expected =
2323            StringArray::from(vec![Some("ccccc"), Some("dd"), Some("eeee"), Some("ccccc")]);
2324        assert_eq!(result.as_string::<i32>(), &expected);
2325
2326        // Nullable path: a null index (position 1) and selected null values
2327        // (sliced indices 0 and 3 are null).
2328        let indices = Int32Array::from(vec![Some(1), None, Some(0), Some(4), Some(3)]);
2329        let result = take(&sliced, &indices, None).unwrap();
2330        let expected = StringArray::from(vec![Some("ccccc"), None, None, Some("eeee"), None]);
2331        assert_eq!(result.as_string::<i32>(), &expected);
2332    }
2333
2334    fn _test_byte_view<T>()
2335    where
2336        T: ByteViewType,
2337        str: AsRef<T::Native>,
2338        T::Native: PartialEq,
2339    {
2340        let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(4), Some(2)]);
2341        let array = {
2342            // ["hello", "world", null, "large payload over 12 bytes", "lulu"]
2343            let mut builder = GenericByteViewBuilder::<T>::new();
2344            builder.append_value("hello");
2345            builder.append_value("world");
2346            builder.append_null();
2347            builder.append_value("large payload over 12 bytes");
2348            builder.append_value("lulu");
2349            builder.finish()
2350        };
2351
2352        let actual = take(&array, &index, None).unwrap();
2353
2354        assert_eq!(actual.len(), index.len());
2355        let actual_buffers = actual.as_byte_view::<T>().data_buffers();
2356        let input_buffers = array.data_buffers();
2357        assert!(Arc::ptr_eq(actual_buffers, input_buffers));
2358
2359        let expected = {
2360            // ["large payload over 12 bytes", null, "world", "large payload over 12 bytes", "lulu", null]
2361            let mut builder = GenericByteViewBuilder::<T>::new();
2362            builder.append_value("large payload over 12 bytes");
2363            builder.append_null();
2364            builder.append_value("world");
2365            builder.append_value("large payload over 12 bytes");
2366            builder.append_value("lulu");
2367            builder.append_null();
2368            builder.finish()
2369        };
2370
2371        assert_eq!(actual.as_ref(), &expected);
2372    }
2373
2374    #[test]
2375    fn test_take_string_view() {
2376        _test_byte_view::<StringViewType>()
2377    }
2378
2379    #[test]
2380    fn test_take_binary_view() {
2381        _test_byte_view::<BinaryViewType>()
2382    }
2383
2384    macro_rules! test_take_list {
2385        ($offset_type:ty, $list_data_type:ident, $list_array_type:ident) => {{
2386            // Construct a value array, [[0,0,0], [-1,-2,-1], [], [2,3]]
2387            let value_data = Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2, 3]).into_data();
2388            // Construct offsets
2389            let value_offsets: [$offset_type; 5] = [0, 3, 6, 6, 8];
2390            let value_offsets = Buffer::from_slice_ref(&value_offsets);
2391            // Construct a list array from the above two
2392            let list_data_type =
2393                DataType::$list_data_type(Arc::new(Field::new_list_field(DataType::Int32, false)));
2394            let list_data = ArrayData::builder(list_data_type.clone())
2395                .len(4)
2396                .add_buffer(value_offsets)
2397                .add_child_data(value_data)
2398                .build()
2399                .unwrap();
2400            let list_array = $list_array_type::from(list_data);
2401
2402            // index returns: [[2,3], null, [-1,-2,-1], [], [0,0,0]]
2403            let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(2), Some(0)]);
2404
2405            let a = take(&list_array, &index, None).unwrap();
2406            let a: &$list_array_type = a.as_any().downcast_ref::<$list_array_type>().unwrap();
2407
2408            // construct a value array with expected results:
2409            // [[2,3], null, [-1,-2,-1], [], [0,0,0]]
2410            let expected_data = Int32Array::from(vec![
2411                Some(2),
2412                Some(3),
2413                Some(-1),
2414                Some(-2),
2415                Some(-1),
2416                Some(0),
2417                Some(0),
2418                Some(0),
2419            ])
2420            .into_data();
2421            // construct offsets
2422            let expected_offsets: [$offset_type; 6] = [0, 2, 2, 5, 5, 8];
2423            let expected_offsets = Buffer::from_slice_ref(&expected_offsets);
2424            // construct list array from the two
2425            let expected_list_data = ArrayData::builder(list_data_type)
2426                .len(5)
2427                // null buffer remains the same as only the indices have nulls
2428                .nulls(index.nulls().cloned())
2429                .add_buffer(expected_offsets)
2430                .add_child_data(expected_data)
2431                .build()
2432                .unwrap();
2433            let expected_list_array = $list_array_type::from(expected_list_data);
2434
2435            assert_eq!(a, &expected_list_array);
2436        }};
2437    }
2438
2439    macro_rules! test_take_list_with_value_nulls {
2440        ($offset_type:ty, $list_data_type:ident, $list_array_type:ident) => {{
2441            // Construct a value array, [[0,null,0], [-1,-2,3], [null], [5,null]]
2442            let value_data = Int32Array::from(vec![
2443                Some(0),
2444                None,
2445                Some(0),
2446                Some(-1),
2447                Some(-2),
2448                Some(3),
2449                None,
2450                Some(5),
2451                None,
2452            ])
2453            .into_data();
2454            // Construct offsets
2455            let value_offsets: [$offset_type; 5] = [0, 3, 6, 7, 9];
2456            let value_offsets = Buffer::from_slice_ref(&value_offsets);
2457            // Construct a list array from the above two
2458            let list_data_type =
2459                DataType::$list_data_type(Arc::new(Field::new_list_field(DataType::Int32, true)));
2460            let list_data = ArrayData::builder(list_data_type.clone())
2461                .len(4)
2462                .add_buffer(value_offsets)
2463                .null_bit_buffer(Some(Buffer::from([0b11111111])))
2464                .add_child_data(value_data)
2465                .build()
2466                .unwrap();
2467            let list_array = $list_array_type::from(list_data);
2468
2469            // index returns: [[null], null, [-1,-2,3], [2,null], [0,null,0]]
2470            let index = UInt32Array::from(vec![Some(2), None, Some(1), Some(3), Some(0)]);
2471
2472            let a = take(&list_array, &index, None).unwrap();
2473            let a: &$list_array_type = a.as_any().downcast_ref::<$list_array_type>().unwrap();
2474
2475            // construct a value array with expected results:
2476            // [[null], null, [-1,-2,3], [5,null], [0,null,0]]
2477            let expected_data = Int32Array::from(vec![
2478                None,
2479                Some(-1),
2480                Some(-2),
2481                Some(3),
2482                Some(5),
2483                None,
2484                Some(0),
2485                None,
2486                Some(0),
2487            ])
2488            .into_data();
2489            // construct offsets
2490            let expected_offsets: [$offset_type; 6] = [0, 1, 1, 4, 6, 9];
2491            let expected_offsets = Buffer::from_slice_ref(&expected_offsets);
2492            // construct list array from the two
2493            let expected_list_data = ArrayData::builder(list_data_type)
2494                .len(5)
2495                // null buffer remains the same as only the indices have nulls
2496                .nulls(index.nulls().cloned())
2497                .add_buffer(expected_offsets)
2498                .add_child_data(expected_data)
2499                .build()
2500                .unwrap();
2501            let expected_list_array = $list_array_type::from(expected_list_data);
2502
2503            assert_eq!(a, &expected_list_array);
2504        }};
2505    }
2506
2507    macro_rules! test_take_list_with_nulls {
2508        ($offset_type:ty, $list_data_type:ident, $list_array_type:ident) => {{
2509            // Construct a value array, [[0,null,0], [-1,-2,3], null, [5,null]]
2510            let value_data = Int32Array::from(vec![
2511                Some(0),
2512                None,
2513                Some(0),
2514                Some(-1),
2515                Some(-2),
2516                Some(3),
2517                Some(5),
2518                None,
2519            ])
2520            .into_data();
2521            // Construct offsets
2522            let value_offsets: [$offset_type; 5] = [0, 3, 6, 6, 8];
2523            let value_offsets = Buffer::from_slice_ref(&value_offsets);
2524            // Construct a list array from the above two
2525            let list_data_type =
2526                DataType::$list_data_type(Arc::new(Field::new_list_field(DataType::Int32, true)));
2527            let list_data = ArrayData::builder(list_data_type.clone())
2528                .len(4)
2529                .add_buffer(value_offsets)
2530                .null_bit_buffer(Some(Buffer::from([0b11111011])))
2531                .add_child_data(value_data)
2532                .build()
2533                .unwrap();
2534            let list_array = $list_array_type::from(list_data);
2535
2536            // index returns: [null, null, [-1,-2,3], [5,null], [0,null,0]]
2537            let index = UInt32Array::from(vec![Some(2), None, Some(1), Some(3), Some(0)]);
2538
2539            let a = take(&list_array, &index, None).unwrap();
2540            let a: &$list_array_type = a.as_any().downcast_ref::<$list_array_type>().unwrap();
2541
2542            // construct a value array with expected results:
2543            // [null, null, [-1,-2,3], [5,null], [0,null,0]]
2544            let expected_data = Int32Array::from(vec![
2545                Some(-1),
2546                Some(-2),
2547                Some(3),
2548                Some(5),
2549                None,
2550                Some(0),
2551                None,
2552                Some(0),
2553            ])
2554            .into_data();
2555            // construct offsets
2556            let expected_offsets: [$offset_type; 6] = [0, 0, 0, 3, 5, 8];
2557            let expected_offsets = Buffer::from_slice_ref(&expected_offsets);
2558            // construct list array from the two
2559            let mut null_bits: [u8; 1] = [0; 1];
2560            bit_util::set_bit(&mut null_bits, 2);
2561            bit_util::set_bit(&mut null_bits, 3);
2562            bit_util::set_bit(&mut null_bits, 4);
2563            let expected_list_data = ArrayData::builder(list_data_type)
2564                .len(5)
2565                // null buffer must be recalculated as both values and indices have nulls
2566                .null_bit_buffer(Some(Buffer::from(null_bits)))
2567                .add_buffer(expected_offsets)
2568                .add_child_data(expected_data)
2569                .build()
2570                .unwrap();
2571            let expected_list_array = $list_array_type::from(expected_list_data);
2572
2573            assert_eq!(a, &expected_list_array);
2574        }};
2575    }
2576
2577    fn test_take_list_view_generic<OffsetType: OffsetSizeTrait, ValuesType: ArrowPrimitiveType, F>(
2578        values: Vec<Option<Vec<Option<ValuesType::Native>>>>,
2579        take_indices: Vec<Option<usize>>,
2580        expected: Vec<Option<Vec<Option<ValuesType::Native>>>>,
2581        mapper: F,
2582    ) where
2583        F: Fn(GenericListViewArray<OffsetType>) -> GenericListViewArray<OffsetType>,
2584    {
2585        let mut list_view_array =
2586            GenericListViewBuilder::<OffsetType, _>::new(PrimitiveBuilder::<ValuesType>::new());
2587
2588        for value in values {
2589            list_view_array.append_option(value);
2590        }
2591        let list_view_array = list_view_array.finish();
2592        let list_view_array = mapper(list_view_array);
2593
2594        let mut indices = UInt64Builder::new();
2595        for idx in take_indices {
2596            indices.append_option(idx.map(|i| i.to_u64().unwrap()));
2597        }
2598        let indices = indices.finish();
2599
2600        let taken = take(&list_view_array, &indices, None)
2601            .unwrap()
2602            .as_list_view()
2603            .clone();
2604
2605        let mut expected_array =
2606            GenericListViewBuilder::<OffsetType, _>::new(PrimitiveBuilder::<ValuesType>::new());
2607        for value in expected {
2608            expected_array.append_option(value);
2609        }
2610        let expected_array = expected_array.finish();
2611
2612        assert_eq!(taken, expected_array);
2613    }
2614
2615    macro_rules! list_view_test_case {
2616        (values: $values:expr, indices: $indices:expr, expected: $expected: expr) => {{
2617            test_take_list_view_generic::<i32, Int8Type, _>($values, $indices, $expected, |x| x);
2618            test_take_list_view_generic::<i64, Int8Type, _>($values, $indices, $expected, |x| x);
2619        }};
2620        (values: $values:expr, transform: $fn:expr, indices: $indices:expr, expected: $expected: expr) => {{
2621            test_take_list_view_generic::<i32, Int8Type, _>($values, $indices, $expected, $fn);
2622            test_take_list_view_generic::<i64, Int8Type, _>($values, $indices, $expected, $fn);
2623        }};
2624    }
2625
2626    fn do_take_fixed_size_list_test<T>(
2627        length: <Int32Type as ArrowPrimitiveType>::Native,
2628        input_data: Vec<Option<Vec<Option<T::Native>>>>,
2629        indices: Vec<<UInt32Type as ArrowPrimitiveType>::Native>,
2630        expected_data: Vec<Option<Vec<Option<T::Native>>>>,
2631    ) where
2632        T: ArrowPrimitiveType,
2633        PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
2634    {
2635        let indices = UInt32Array::from(indices);
2636
2637        let input_array = FixedSizeListArray::from_iter_primitive::<T, _, _>(input_data, length);
2638
2639        let output =
2640            take_fixed_size_list::<_, true>(&input_array, &indices, length as u32).unwrap();
2641
2642        let expected = FixedSizeListArray::from_iter_primitive::<T, _, _>(expected_data, length);
2643
2644        assert_eq!(&output, &expected)
2645    }
2646
2647    #[test]
2648    // Fast path (primitive child, no child nulls) with null indices — verifies offset backfill for null slots.
2649    fn test_take_list_primitive_child_null_indices() {
2650        // Row sizes deliberately vary (1, 3, 2) so the test exercises
2651        // non-uniform offset arithmetic, not just uniform stride.
2652        let list = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
2653            Some(vec![Some(1)]),
2654            Some(vec![Some(2), Some(3), Some(4)]),
2655            Some(vec![Some(5), Some(6)]),
2656        ]);
2657        let indices = Int32Array::from(vec![Some(2), None, Some(0), Some(1)]);
2658        let result = take(&list, &indices, None).unwrap();
2659        let result = result.as_any().downcast_ref::<ListArray>().unwrap();
2660
2661        let expected = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
2662            Some(vec![Some(5), Some(6)]),
2663            None,
2664            Some(vec![Some(1)]),
2665            Some(vec![Some(2), Some(3), Some(4)]),
2666        ]);
2667        assert_eq!(result, &expected);
2668    }
2669
2670    #[test]
2671    fn test_take_list() {
2672        test_take_list!(i32, List, ListArray);
2673    }
2674
2675    #[test]
2676    fn test_take_large_list() {
2677        test_take_list!(i64, LargeList, LargeListArray);
2678    }
2679
2680    #[test]
2681    fn test_take_list_with_value_nulls() {
2682        test_take_list_with_value_nulls!(i32, List, ListArray);
2683    }
2684
2685    #[test]
2686    fn test_take_large_list_with_value_nulls() {
2687        test_take_list_with_value_nulls!(i64, LargeList, LargeListArray);
2688    }
2689
2690    #[test]
2691    fn test_test_take_list_with_nulls() {
2692        test_take_list_with_nulls!(i32, List, ListArray);
2693    }
2694
2695    #[test]
2696    fn test_test_take_large_list_with_nulls() {
2697        test_take_list_with_nulls!(i64, LargeList, LargeListArray);
2698    }
2699
2700    #[test]
2701    fn test_test_take_list_view_reversed() {
2702        // Take reversed indices
2703        list_view_test_case! {
2704            values: vec![
2705                Some(vec![Some(1), None, Some(3)]),
2706                None,
2707                Some(vec![Some(7), Some(8), None]),
2708            ],
2709            indices: vec![Some(2), Some(1), Some(0)],
2710            expected: vec![
2711                Some(vec![Some(7), Some(8), None]),
2712                None,
2713                Some(vec![Some(1), None, Some(3)]),
2714            ]
2715        }
2716    }
2717
2718    #[test]
2719    fn test_take_list_view_null_indices() {
2720        // Take with null indices
2721        list_view_test_case! {
2722            values: vec![
2723                Some(vec![Some(1), None, Some(3)]),
2724                None,
2725                Some(vec![Some(7), Some(8), None]),
2726            ],
2727            indices: vec![None, Some(0), None],
2728            expected: vec![None, Some(vec![Some(1), None, Some(3)]), None]
2729        }
2730    }
2731
2732    #[test]
2733    fn test_take_list_view_null_values() {
2734        // Take at null values
2735        list_view_test_case! {
2736            values: vec![
2737                Some(vec![Some(1), None, Some(3)]),
2738                None,
2739                Some(vec![Some(7), Some(8), None]),
2740            ],
2741            indices: vec![Some(1), Some(1), Some(1), None, None],
2742            expected: vec![None; 5]
2743        }
2744    }
2745
2746    #[test]
2747    fn test_take_list_view_sliced() {
2748        // Take null indices/values, with slicing.
2749        list_view_test_case! {
2750            values: vec![
2751                Some(vec![Some(1)]),
2752                None,
2753                None,
2754                Some(vec![Some(2), Some(3)]),
2755                Some(vec![Some(4), Some(5)]),
2756                None,
2757            ],
2758            transform: |l| l.slice(2, 4),
2759            indices: vec![Some(0), Some(3), None, Some(1), Some(2)],
2760            expected: vec![
2761                None, None, None, Some(vec![Some(2), Some(3)]), Some(vec![Some(4), Some(5)])
2762            ]
2763        }
2764    }
2765
2766    #[test]
2767    fn test_take_fixed_size_list() {
2768        do_take_fixed_size_list_test::<Int32Type>(
2769            3,
2770            vec![
2771                Some(vec![None, Some(1), Some(2)]),
2772                Some(vec![Some(3), Some(4), None]),
2773                Some(vec![Some(6), Some(7), Some(8)]),
2774            ],
2775            vec![2, 1, 0],
2776            vec![
2777                Some(vec![Some(6), Some(7), Some(8)]),
2778                Some(vec![Some(3), Some(4), None]),
2779                Some(vec![None, Some(1), Some(2)]),
2780            ],
2781        );
2782
2783        do_take_fixed_size_list_test::<UInt8Type>(
2784            1,
2785            vec![
2786                Some(vec![Some(1)]),
2787                Some(vec![Some(2)]),
2788                Some(vec![Some(3)]),
2789                Some(vec![Some(4)]),
2790                Some(vec![Some(5)]),
2791                Some(vec![Some(6)]),
2792                Some(vec![Some(7)]),
2793                Some(vec![Some(8)]),
2794            ],
2795            vec![2, 7, 0],
2796            vec![
2797                Some(vec![Some(3)]),
2798                Some(vec![Some(8)]),
2799                Some(vec![Some(1)]),
2800            ],
2801        );
2802
2803        do_take_fixed_size_list_test::<UInt64Type>(
2804            3,
2805            vec![
2806                Some(vec![Some(10), Some(11), Some(12)]),
2807                Some(vec![Some(13), Some(14), Some(15)]),
2808                None,
2809                Some(vec![Some(16), Some(17), Some(18)]),
2810            ],
2811            vec![3, 2, 1, 2, 0],
2812            vec![
2813                Some(vec![Some(16), Some(17), Some(18)]),
2814                None,
2815                Some(vec![Some(13), Some(14), Some(15)]),
2816                None,
2817                Some(vec![Some(10), Some(11), Some(12)]),
2818            ],
2819        );
2820    }
2821
2822    #[test]
2823    fn test_take_fixed_size_binary_with_nulls_indices() {
2824        let fsb = FixedSizeBinaryArray::try_from_sparse_iter_with_size(
2825            [
2826                Some(vec![0x01, 0x01, 0x01, 0x01]),
2827                Some(vec![0x02, 0x02, 0x02, 0x02]),
2828                Some(vec![0x03, 0x03, 0x03, 0x03]),
2829                Some(vec![0x04, 0x04, 0x04, 0x04]),
2830            ]
2831            .into_iter(),
2832            4,
2833        )
2834        .unwrap();
2835
2836        // The two middle indices are null -> Should be null in the output.
2837        let indices = UInt32Array::from(vec![Some(0), None, None, Some(3)]);
2838
2839        let result = take_fixed_size_binary::<_, true>(&fsb, &indices, 4).unwrap();
2840        assert_eq!(result.len(), 4);
2841        assert_eq!(result.null_count(), 2);
2842        assert_eq!(
2843            result.nulls().unwrap().iter().collect::<Vec<_>>(),
2844            vec![true, false, false, true]
2845        );
2846    }
2847
2848    /// The [`take_fixed_size_binary`] kernel contains optimizations that provide a faster
2849    /// implementation for commonly-used value lengths. This test uses a value length that is not
2850    /// optimized to test both code paths.
2851    #[test]
2852    fn test_take_fixed_size_binary_with_nulls_indices_not_optimized_length() {
2853        let fsb = FixedSizeBinaryArray::try_from_sparse_iter_with_size(
2854            [
2855                Some(vec![0x01, 0x01, 0x01, 0x01, 0x01]),
2856                Some(vec![0x02, 0x02, 0x02, 0x02, 0x01]),
2857                Some(vec![0x03, 0x03, 0x03, 0x03, 0x01]),
2858                Some(vec![0x04, 0x04, 0x04, 0x04, 0x01]),
2859            ]
2860            .into_iter(),
2861            5,
2862        )
2863        .unwrap();
2864
2865        // The two middle indices are null -> Should be null in the output.
2866        let indices = UInt32Array::from(vec![Some(0), None, None, Some(3)]);
2867
2868        let result = take_fixed_size_binary::<_, true>(&fsb, &indices, 5).unwrap();
2869        assert_eq!(result.len(), 4);
2870        assert_eq!(result.null_count(), 2);
2871        assert_eq!(
2872            result.nulls().unwrap().iter().collect::<Vec<_>>(),
2873            vec![true, false, false, true]
2874        );
2875    }
2876
2877    #[test]
2878    #[should_panic(expected = "index out of bounds: the len is 4 but the index is 1000")]
2879    fn test_take_list_out_of_bounds() {
2880        // Construct a value array, [[0,0,0], [-1,-2,-1], [2,3]]
2881        let value_data = Int32Array::from(vec![0, 0, 0, -1, -2, -1, 2, 3]).into_data();
2882        // Construct offsets
2883        let value_offsets = Buffer::from_slice_ref([0, 3, 6, 8]);
2884        // Construct a list array from the above two
2885        let list_data_type =
2886            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
2887        let list_data = ArrayData::builder(list_data_type)
2888            .len(3)
2889            .add_buffer(value_offsets)
2890            .add_child_data(value_data)
2891            .build()
2892            .unwrap();
2893        let list_array = ListArray::from(list_data);
2894
2895        let index = UInt32Array::from(vec![1000]);
2896
2897        // A panic is expected here since we have not supplied the check_bounds
2898        // option.
2899        take(&list_array, &index, None).unwrap();
2900    }
2901
2902    #[test]
2903    fn test_take_map() {
2904        let values = Int32Array::from(vec![1, 2, 3, 4]);
2905        let array =
2906            MapArray::new_from_strings(vec!["a", "b", "c", "a"].into_iter(), &values, &[0, 3, 4])
2907                .unwrap();
2908
2909        let index = UInt32Array::from(vec![0]);
2910
2911        let result = take(&array, &index, None).unwrap();
2912        let expected: ArrayRef = Arc::new(
2913            MapArray::new_from_strings(
2914                vec!["a", "b", "c"].into_iter(),
2915                &values.slice(0, 3),
2916                &[0, 3],
2917            )
2918            .unwrap(),
2919        );
2920        assert_eq!(&expected, &result);
2921    }
2922
2923    #[test]
2924    fn test_take_struct() {
2925        let array = create_test_struct(vec![
2926            Some((Some(true), Some(42))),
2927            Some((Some(false), Some(28))),
2928            Some((Some(false), Some(19))),
2929            Some((Some(true), Some(31))),
2930            None,
2931        ]);
2932
2933        let index = UInt32Array::from(vec![0, 3, 1, 0, 2, 4]);
2934        let actual = take(&array, &index, None).unwrap();
2935        let actual: &StructArray = actual.as_any().downcast_ref::<StructArray>().unwrap();
2936        assert_eq!(index.len(), actual.len());
2937        assert_eq!(1, actual.null_count());
2938
2939        let expected = create_test_struct(vec![
2940            Some((Some(true), Some(42))),
2941            Some((Some(true), Some(31))),
2942            Some((Some(false), Some(28))),
2943            Some((Some(true), Some(42))),
2944            Some((Some(false), Some(19))),
2945            None,
2946        ]);
2947
2948        assert_eq!(&expected, actual);
2949
2950        let nulls = NullBuffer::from(&[false, true, false, true, false, true]);
2951        let empty_struct_arr = StructArray::new_empty_fields(6, Some(nulls));
2952        let index = UInt32Array::from(vec![0, 2, 1, 4]);
2953        let actual = take(&empty_struct_arr, &index, None).unwrap();
2954
2955        let expected_nulls = NullBuffer::from(&[false, false, true, false]);
2956        let expected_struct_arr = StructArray::new_empty_fields(4, Some(expected_nulls));
2957        assert_eq!(&expected_struct_arr, actual.as_struct());
2958    }
2959
2960    #[test]
2961    fn test_take_struct_with_null_indices() {
2962        let array = create_test_struct(vec![
2963            Some((Some(true), Some(42))),
2964            Some((Some(false), Some(28))),
2965            Some((Some(false), Some(19))),
2966            Some((Some(true), Some(31))),
2967            None,
2968        ]);
2969
2970        let index = UInt32Array::from(vec![None, Some(3), Some(1), None, Some(0), Some(4)]);
2971        let actual = take(&array, &index, None).unwrap();
2972        let actual: &StructArray = actual.as_any().downcast_ref::<StructArray>().unwrap();
2973        assert_eq!(index.len(), actual.len());
2974        assert_eq!(3, actual.null_count()); // 2 because of indices, 1 because of struct array
2975
2976        let expected = create_test_struct(vec![
2977            None,
2978            Some((Some(true), Some(31))),
2979            Some((Some(false), Some(28))),
2980            None,
2981            Some((Some(true), Some(42))),
2982            None,
2983        ]);
2984
2985        assert_eq!(&expected, actual);
2986    }
2987
2988    #[test]
2989    fn test_take_out_of_bounds() {
2990        let index = UInt32Array::from(vec![Some(3), None, Some(1), Some(3), Some(6)]);
2991        let take_opt = TakeOptions { check_bounds: true };
2992
2993        // int64
2994        let result = test_take_primitive_arrays::<Int64Type>(
2995            vec![Some(0), None, Some(2), Some(3), None],
2996            &index,
2997            Some(take_opt),
2998            vec![None],
2999        );
3000        assert!(result.is_err());
3001    }
3002
3003    #[test]
3004    #[should_panic(expected = "index out of bounds: the len is 4 but the index is 1000")]
3005    fn test_take_out_of_bounds_panic() {
3006        let index = UInt32Array::from(vec![Some(1000)]);
3007
3008        test_take_primitive_arrays::<Int64Type>(
3009            vec![Some(0), Some(1), Some(2), Some(3)],
3010            &index,
3011            None,
3012            vec![None],
3013        )
3014        .unwrap();
3015    }
3016
3017    #[test]
3018    fn test_null_array_smaller_than_indices() {
3019        let values = NullArray::new(2);
3020        let indices = UInt32Array::from(vec![Some(0), None, Some(15)]);
3021
3022        let result = take(&values, &indices, None).unwrap();
3023        let expected: ArrayRef = Arc::new(NullArray::new(3));
3024        assert_eq!(&result, &expected);
3025    }
3026
3027    #[test]
3028    fn test_null_array_larger_than_indices() {
3029        let values = NullArray::new(5);
3030        let indices = UInt32Array::from(vec![Some(0), None, Some(15)]);
3031
3032        let result = take(&values, &indices, None).unwrap();
3033        let expected: ArrayRef = Arc::new(NullArray::new(3));
3034        assert_eq!(&result, &expected);
3035    }
3036
3037    #[test]
3038    fn test_null_array_indices_out_of_bounds() {
3039        let values = NullArray::new(5);
3040        let indices = UInt32Array::from(vec![Some(0), None, Some(15)]);
3041
3042        let result = take(&values, &indices, Some(TakeOptions { check_bounds: true }));
3043        assert_eq!(
3044            result.unwrap_err().to_string(),
3045            "Compute error: Array index out of bounds, cannot get item at index 15 from 5 entries"
3046        );
3047    }
3048
3049    #[test]
3050    fn test_take_dict() {
3051        let mut dict_builder = StringDictionaryBuilder::<Int16Type>::new();
3052
3053        dict_builder.append("foo").unwrap();
3054        dict_builder.append("bar").unwrap();
3055        dict_builder.append("").unwrap();
3056        dict_builder.append_null();
3057        dict_builder.append("foo").unwrap();
3058        dict_builder.append("bar").unwrap();
3059        dict_builder.append("bar").unwrap();
3060        dict_builder.append("foo").unwrap();
3061
3062        let array = dict_builder.finish();
3063        let dict_values = array.values().clone();
3064        let dict_values = dict_values.as_any().downcast_ref::<StringArray>().unwrap();
3065
3066        let indices = UInt32Array::from(vec![
3067            Some(0), // first "foo"
3068            Some(7), // last "foo"
3069            None,    // null index should return null
3070            Some(5), // second "bar"
3071            Some(6), // another "bar"
3072            Some(2), // empty string
3073            Some(3), // input is null at this index
3074        ]);
3075
3076        let result = take(&array, &indices, None).unwrap();
3077        let result = result
3078            .as_any()
3079            .downcast_ref::<DictionaryArray<Int16Type>>()
3080            .unwrap();
3081
3082        let result_values: StringArray = result.values().to_data().into();
3083
3084        // dictionary values should stay the same
3085        let expected_values = StringArray::from(vec!["foo", "bar", ""]);
3086        assert_eq!(&expected_values, dict_values);
3087        assert_eq!(&expected_values, &result_values);
3088
3089        let expected_keys = Int16Array::from(vec![
3090            Some(0),
3091            Some(0),
3092            None,
3093            Some(1),
3094            Some(1),
3095            Some(2),
3096            None,
3097        ]);
3098        assert_eq!(result.keys(), &expected_keys);
3099    }
3100
3101    fn build_generic_list<S, T>(data: Vec<Option<Vec<T::Native>>>) -> GenericListArray<S>
3102    where
3103        S: OffsetSizeTrait + 'static,
3104        T: ArrowPrimitiveType,
3105        PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
3106    {
3107        GenericListArray::from_iter_primitive::<T, _, _>(
3108            data.iter()
3109                .map(|x| x.as_ref().map(|x| x.iter().map(|x| Some(*x)))),
3110        )
3111    }
3112
3113    fn test_take_sliced_list_generic<S: OffsetSizeTrait + 'static>() {
3114        let list = build_generic_list::<S, Int32Type>(vec![
3115            Some(vec![0, 1]),
3116            Some(vec![2, 3, 4]),
3117            None,
3118            Some(vec![]),
3119            Some(vec![5, 6]),
3120            Some(vec![7]),
3121        ]);
3122        let sliced = list.slice(1, 4);
3123        let indices = UInt32Array::from(vec![Some(3), Some(0), None, Some(2), Some(1)]);
3124
3125        let taken = take(&sliced, &indices, None).unwrap();
3126        let taken = taken.as_list::<S>();
3127
3128        let expected = build_generic_list::<S, Int32Type>(vec![
3129            Some(vec![5, 6]),
3130            Some(vec![2, 3, 4]),
3131            None,
3132            Some(vec![]),
3133            None,
3134        ]);
3135
3136        assert_eq!(taken, &expected);
3137    }
3138
3139    fn test_take_sliced_list_with_value_nulls_generic<S: OffsetSizeTrait + 'static>() {
3140        let list = GenericListArray::<S>::from_iter_primitive::<Int32Type, _, _>(vec![
3141            Some(vec![Some(10)]),
3142            Some(vec![None, Some(1)]),
3143            None,
3144            Some(vec![Some(2), None]),
3145            Some(vec![]),
3146            Some(vec![Some(3)]),
3147        ]);
3148        let sliced = list.slice(1, 4);
3149        let indices = UInt32Array::from(vec![Some(2), Some(0), None, Some(3), Some(1)]);
3150
3151        let taken = take(&sliced, &indices, None).unwrap();
3152        let taken = taken.as_list::<S>();
3153
3154        let expected = GenericListArray::<S>::from_iter_primitive::<Int32Type, _, _>(vec![
3155            Some(vec![Some(2), None]),
3156            Some(vec![None, Some(1)]),
3157            None,
3158            Some(vec![]),
3159            None,
3160        ]);
3161
3162        assert_eq!(taken, &expected);
3163    }
3164
3165    #[test]
3166    fn test_take_sliced_list() {
3167        test_take_sliced_list_generic::<i32>();
3168    }
3169
3170    #[test]
3171    fn test_take_sliced_large_list() {
3172        test_take_sliced_list_generic::<i64>();
3173    }
3174
3175    #[test]
3176    fn test_take_sliced_list_with_value_nulls() {
3177        test_take_sliced_list_with_value_nulls_generic::<i32>();
3178    }
3179
3180    #[test]
3181    fn test_take_sliced_large_list_with_value_nulls() {
3182        test_take_sliced_list_with_value_nulls_generic::<i64>();
3183    }
3184
3185    #[test]
3186    fn test_take_runs() {
3187        let logical_array: Vec<i32> = vec![1_i32, 1, 2, 2, 1, 1, 1, 2, 2, 1, 1, 2, 2];
3188
3189        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
3190        builder.extend(logical_array.into_iter().map(Some));
3191        let run_array = builder.finish();
3192
3193        let take_indices: PrimitiveArray<Int32Type> =
3194            vec![7, 2, 3, 7, 11, 4, 6].into_iter().collect();
3195
3196        let take_out = take_run(&run_array, &take_indices).unwrap();
3197
3198        assert_eq!(take_out.len(), 7);
3199        // adjacent identical values are merged: [2,2,2,2,2,1,1] -> 2 runs
3200        assert_eq!(
3201            take_out.run_ends().values().len(),
3202            2,
3203            "expected two physical runs"
3204        );
3205        assert_eq!(take_out.run_ends().values(), &[5_i32, 7]);
3206
3207        let take_out_values = take_out.values().as_primitive::<Int32Type>();
3208        assert_eq!(take_out_values.values(), &[2, 1]);
3209    }
3210
3211    #[test]
3212    fn test_take_runs_null_indices() {
3213        // A null index must not become logical index 0, and null indices must merge with
3214        // consecutive runs whose values are already null.
3215        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
3216        builder.extend([Some(10), Some(10), None, None, Some(99)]);
3217        let run_array = builder.finish();
3218
3219        let indices = Int32Array::from(vec![Some(0), None, Some(2), Some(3), Some(4)]);
3220        let taken = take(&run_array, &indices, None).unwrap();
3221        let run = taken.as_run::<Int32Type>();
3222        let logical: Vec<Option<i32>> = run.downcast::<Int32Array>().unwrap().into_iter().collect();
3223        assert_eq!(logical, vec![Some(10), None, None, None, Some(99)]);
3224        assert_eq!(run.run_ends().values(), &[1_i32, 4, 5]);
3225    }
3226
3227    #[test]
3228    fn test_take_runs_sliced() {
3229        let logical_array: Vec<i32> = vec![1, 1, 2, 2, 3, 3, 3, 4, 4, 5, 5, 6, 6];
3230
3231        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
3232        builder.extend(logical_array.into_iter().map(Some));
3233        let run_array = builder.finish();
3234
3235        let run_array = run_array.slice(4, 6); // [3, 3, 3, 4, 4, 5]
3236
3237        let take_indices: PrimitiveArray<Int32Type> = vec![0, 5, 5, 1, 4].into_iter().collect();
3238
3239        let result = take_run(&run_array, &take_indices).unwrap();
3240        let result = result.downcast::<Int32Array>().unwrap();
3241
3242        // [3, 5, 5, 3, 4] -> 4 physical runs (no adjacent duplicates to merge)
3243        assert_eq!(
3244            result.run_ends().values().len(),
3245            4,
3246            "expected four physical runs"
3247        );
3248        assert_eq!(result.run_ends().values(), &[1_i32, 3, 4, 5]);
3249
3250        let expected = vec![3, 5, 5, 3, 4];
3251        let actual = result.into_iter().flatten().collect::<Vec<_>>();
3252
3253        assert_eq!(expected, actual);
3254    }
3255
3256    #[test]
3257    fn test_take_value_index_from_fixed_list() {
3258        let list = FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>(
3259            vec![
3260                Some(vec![Some(1), Some(2), None]),
3261                Some(vec![Some(4), None, Some(6)]),
3262                None,
3263                Some(vec![None, Some(8), Some(9)]),
3264            ],
3265            3,
3266        );
3267
3268        let indices = UInt32Array::from(vec![2, 1, 0]);
3269        let indexed = take_value_indices_from_fixed_size_list(&list, &indices, 3).unwrap();
3270
3271        assert_eq!(indexed, UInt32Array::from(vec![6, 7, 8, 3, 4, 5, 0, 1, 2]));
3272
3273        let indices = UInt32Array::from(vec![3, 2, 1, 2, 0]);
3274        let indexed = take_value_indices_from_fixed_size_list(&list, &indices, 3).unwrap();
3275
3276        assert_eq!(
3277            indexed,
3278            UInt32Array::from(vec![9, 10, 11, 6, 7, 8, 3, 4, 5, 6, 7, 8, 0, 1, 2])
3279        );
3280    }
3281
3282    #[test]
3283    fn test_take_null_indices() {
3284        // Build indices with values that are out of bounds, but masked by null mask
3285        let indices = Int32Array::new(
3286            vec![1, 2, 400, 400].into(),
3287            Some(NullBuffer::from(vec![true, true, false, false])),
3288        );
3289        let values = Int32Array::from(vec![1, 23, 4, 5]);
3290        let r = take(&values, &indices, None).unwrap();
3291        let values = r
3292            .as_primitive::<Int32Type>()
3293            .into_iter()
3294            .collect::<Vec<_>>();
3295        assert_eq!(&values, &[Some(23), Some(4), None, None])
3296    }
3297
3298    #[test]
3299    fn test_take_fixed_size_list_null_indices() {
3300        let indices = Int32Array::from_iter([Some(0), None]);
3301        let values = Arc::new(Int32Array::from(vec![0, 1, 2, 3]));
3302        let arr_field = Arc::new(Field::new_list_field(values.data_type().clone(), true));
3303        let values = FixedSizeListArray::try_new(arr_field, 2, values, None).unwrap();
3304
3305        let r = take(&values, &indices, None).unwrap();
3306        let values = r
3307            .as_fixed_size_list()
3308            .values()
3309            .as_primitive::<Int32Type>()
3310            .into_iter()
3311            .collect::<Vec<_>>();
3312        assert_eq!(values, &[Some(0), Some(1), None, None])
3313    }
3314
3315    #[test]
3316    fn test_take_bytes_null_indices() {
3317        let indices = Int32Array::new(
3318            vec![0, 1, 400, 400].into(),
3319            Some(NullBuffer::from_iter(vec![true, true, false, false])),
3320        );
3321        let values = StringArray::from(vec![Some("foo"), None]);
3322        let r = take(&values, &indices, None).unwrap();
3323        let values = r.as_string::<i32>().iter().collect::<Vec<_>>();
3324        assert_eq!(&values, &[Some("foo"), None, None, None])
3325    }
3326
3327    #[test]
3328    fn test_take_union_sparse() {
3329        let structs = create_test_struct(vec![
3330            Some((Some(true), Some(42))),
3331            Some((Some(false), Some(28))),
3332            Some((Some(false), Some(19))),
3333            Some((Some(true), Some(31))),
3334            None,
3335        ]);
3336        let strings = StringArray::from(vec![Some("a"), None, Some("c"), None, Some("d")]);
3337        let type_ids = [1; 5].into_iter().collect::<ScalarBuffer<i8>>();
3338
3339        let union_fields = [
3340            (
3341                0,
3342                Arc::new(Field::new("f1", structs.data_type().clone(), true)),
3343            ),
3344            (
3345                1,
3346                Arc::new(Field::new("f2", strings.data_type().clone(), true)),
3347            ),
3348        ]
3349        .into_iter()
3350        .collect();
3351        let children = vec![Arc::new(structs) as Arc<dyn Array>, Arc::new(strings)];
3352        let array = UnionArray::try_new(union_fields, type_ids, None, children).unwrap();
3353
3354        let indices = vec![0, 3, 1, 0, 2, 4];
3355        let index = UInt32Array::from(indices.clone());
3356        let actual = take(&array, &index, None).unwrap();
3357        let actual = actual.as_any().downcast_ref::<UnionArray>().unwrap();
3358        let strings = actual.child(1);
3359        let strings = strings.as_any().downcast_ref::<StringArray>().unwrap();
3360
3361        let actual = strings.iter().collect::<Vec<_>>();
3362        let expected = vec![Some("a"), None, None, Some("a"), Some("c"), Some("d")];
3363        assert_eq!(expected, actual);
3364    }
3365
3366    #[test]
3367    fn test_take_union_dense() {
3368        let type_ids = vec![0, 1, 1, 0, 0, 1, 0];
3369        let offsets = vec![0, 0, 1, 1, 2, 2, 3];
3370        let ints = vec![10, 20, 30, 40];
3371        let strings = vec![Some("a"), None, Some("c"), Some("d")];
3372
3373        let indices = vec![0, 3, 1, 0, 2, 4];
3374
3375        let taken_type_ids = vec![0, 0, 1, 0, 1, 0];
3376        let taken_offsets = vec![0, 1, 0, 2, 1, 3];
3377        let taken_ints = vec![10, 20, 10, 30];
3378        let taken_strings = vec![Some("a"), None];
3379
3380        let type_ids = <ScalarBuffer<i8>>::from(type_ids);
3381        let offsets = <ScalarBuffer<i32>>::from(offsets);
3382        let ints = UInt32Array::from(ints);
3383        let strings = StringArray::from(strings);
3384
3385        let union_fields = [
3386            (
3387                0,
3388                Arc::new(Field::new("f1", ints.data_type().clone(), true)),
3389            ),
3390            (
3391                1,
3392                Arc::new(Field::new("f2", strings.data_type().clone(), true)),
3393            ),
3394        ]
3395        .into_iter()
3396        .collect();
3397
3398        let array = UnionArray::try_new(
3399            union_fields,
3400            type_ids,
3401            Some(offsets),
3402            vec![Arc::new(ints), Arc::new(strings)],
3403        )
3404        .unwrap();
3405
3406        let index = UInt32Array::from(indices);
3407
3408        let actual = take(&array, &index, None).unwrap();
3409        let actual = actual.as_any().downcast_ref::<UnionArray>().unwrap();
3410
3411        assert_eq!(actual.offsets(), Some(&ScalarBuffer::from(taken_offsets)));
3412        assert_eq!(actual.type_ids(), &ScalarBuffer::from(taken_type_ids));
3413        assert_eq!(
3414            UInt32Array::from(actual.child(0).to_data()),
3415            UInt32Array::from(taken_ints)
3416        );
3417        assert_eq!(
3418            StringArray::from(actual.child(1).to_data()),
3419            StringArray::from(taken_strings)
3420        );
3421    }
3422
3423    fn union_i32_logical(array: &UnionArray) -> Vec<Option<i32>> {
3424        (0..array.len())
3425            .map(|i| {
3426                let child = array.child(array.type_id(i)).as_primitive::<Int32Type>();
3427                let offset = array.value_offset(i);
3428                if child.is_null(offset) {
3429                    None
3430                } else {
3431                    Some(child.value(offset))
3432                }
3433            })
3434            .collect()
3435    }
3436
3437    #[test]
3438    fn test_take_union_dense_null_indices() {
3439        // Dense [1, 2, 3]; a null index must not become child offset 0. Use a non-zero
3440        // type id to verify null indices don't use the invalid default type id 0.
3441        let fields =
3442            UnionFields::try_new(vec![5], vec![Field::new("i", DataType::Int32, true)]).unwrap();
3443        let dense = UnionArray::try_new(
3444            fields.clone(),
3445            ScalarBuffer::from(vec![5_i8, 5, 5]),
3446            Some(ScalarBuffer::from(vec![0_i32, 1, 2])),
3447            vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
3448        )
3449        .unwrap();
3450        let sparse = UnionArray::try_new(
3451            fields,
3452            ScalarBuffer::from(vec![5_i8, 5, 5]),
3453            None,
3454            vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
3455        )
3456        .unwrap();
3457
3458        // Use an out-of-bounds backing value for the null index to exercise `take_native`'s
3459        // default type id path. A union need not have a child with type id 0.
3460        let indices = UInt32Array::new(
3461            ScalarBuffer::from(vec![0_u32, 99, 2]),
3462            Some(NullBuffer::from(vec![true, false, true])),
3463        );
3464        let dense_taken = take(&dense, &indices, None).unwrap();
3465        let sparse_taken = take(&sparse, &indices, None).unwrap();
3466        let dense_logical = union_i32_logical(dense_taken.as_any().downcast_ref().unwrap());
3467        let sparse_logical = union_i32_logical(sparse_taken.as_any().downcast_ref().unwrap());
3468
3469        assert_eq!(dense_logical, vec![Some(1), None, Some(3)]);
3470        assert_eq!(dense_logical, sparse_logical);
3471    }
3472
3473    #[test]
3474    fn test_take_empty_union_without_null_indices() {
3475        let fields = UnionFields::try_new(vec![], Vec::<Field>::new()).unwrap();
3476        let indices = UInt32Array::from(Vec::<u32>::new());
3477
3478        let sparse = UnionArray::try_new(
3479            fields.clone(),
3480            ScalarBuffer::<i8>::from(vec![]),
3481            None,
3482            vec![],
3483        )
3484        .unwrap();
3485        let dense = UnionArray::try_new(
3486            fields,
3487            ScalarBuffer::<i8>::from(vec![]),
3488            Some(ScalarBuffer::<i32>::from(vec![])),
3489            vec![],
3490        )
3491        .unwrap();
3492
3493        for values in [&sparse, &dense] {
3494            let taken = take(values, &indices, None).unwrap();
3495            assert_eq!(taken.len(), 0);
3496            assert_eq!(taken.data_type(), values.data_type());
3497        }
3498    }
3499
3500    #[test]
3501    fn test_take_empty_union_with_null_indices() {
3502        let fields = UnionFields::try_new(vec![], Vec::<Field>::new()).unwrap();
3503        let indices = UInt32Array::from(vec![None]);
3504
3505        let sparse = UnionArray::try_new(
3506            fields.clone(),
3507            ScalarBuffer::<i8>::from(vec![]),
3508            None,
3509            vec![],
3510        )
3511        .unwrap();
3512        let dense = UnionArray::try_new(
3513            fields,
3514            ScalarBuffer::<i8>::from(vec![]),
3515            Some(ScalarBuffer::<i32>::from(vec![])),
3516            vec![],
3517        )
3518        .unwrap();
3519
3520        for values in [&sparse, &dense] {
3521            let error = take(values, &indices, None).unwrap_err();
3522            assert_eq!(
3523                error.to_string(),
3524                "Compute error: Cannot take from a union with zero fields when indices contains nulls"
3525            );
3526        }
3527    }
3528
3529    #[test]
3530    fn test_take_union_dense_using_builder() {
3531        let mut builder = UnionBuilder::new_dense();
3532
3533        builder.append::<Int32Type>("a", 1).unwrap();
3534        builder.append::<Float64Type>("b", 3.0).unwrap();
3535        builder.append::<Int32Type>("a", 4).unwrap();
3536        builder.append::<Int32Type>("a", 5).unwrap();
3537        builder.append::<Float64Type>("b", 2.0).unwrap();
3538
3539        let union = builder.build().unwrap();
3540
3541        let indices = UInt32Array::from(vec![2, 0, 1, 2]);
3542
3543        let mut builder = UnionBuilder::new_dense();
3544
3545        builder.append::<Int32Type>("a", 4).unwrap();
3546        builder.append::<Int32Type>("a", 1).unwrap();
3547        builder.append::<Float64Type>("b", 3.0).unwrap();
3548        builder.append::<Int32Type>("a", 4).unwrap();
3549
3550        let taken = builder.build().unwrap();
3551
3552        assert_eq!(
3553            taken.to_data(),
3554            take(&union, &indices, None).unwrap().to_data()
3555        );
3556    }
3557
3558    #[test]
3559    fn test_take_union_dense_all_match_issue_6206() {
3560        let fields = UnionFields::from_fields(vec![Field::new("a", DataType::Int64, false)]);
3561        let ints = Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5]));
3562
3563        let array = UnionArray::try_new(
3564            fields,
3565            ScalarBuffer::from(vec![0_i8, 0, 0, 0, 0]),
3566            Some(ScalarBuffer::from_iter(0_i32..5)),
3567            vec![ints],
3568        )
3569        .unwrap();
3570
3571        let indices = Int64Array::from(vec![0, 2, 4]);
3572        let array = take(&array, &indices, None).unwrap();
3573        assert_eq!(array.len(), 3);
3574    }
3575
3576    /// Fixture for the offset-overflow tests: a single large value plus the
3577    /// number of times it must be selected so the cumulative offset exceeds
3578    /// `i32::MAX`. Using a large value keeps the index count (and the test
3579    /// runtime) small.
3580    fn offset_overflow_fixture() -> (StringArray, usize) {
3581        let value_len = 1_000_000;
3582        let values = StringArray::from(vec![Some("a".repeat(value_len))]);
3583        let n = i32::MAX as usize / value_len + 1;
3584        (values, n)
3585    }
3586
3587    #[test]
3588    fn test_take_bytes_offset_overflow() {
3589        let (values, n) = offset_overflow_fixture();
3590        let indices = Int32Array::from(vec![0; n]);
3591        assert!(matches!(
3592            take(&values, &indices, None),
3593            Err(ArrowError::OffsetOverflowError(_))
3594        ));
3595    }
3596
3597    /// The offset-overflow error must also be produced on the nullable code
3598    /// path (when the output contains nulls), not only on the no-null fast path.
3599    #[test]
3600    fn test_take_bytes_offset_overflow_nullable() {
3601        let (values, n) = offset_overflow_fixture();
3602        // A null index forces the output to contain nulls, exercising the
3603        // nullable code path.
3604        let validity =
3605            NullBuffer::from_iter(std::iter::once(false).chain(std::iter::repeat_n(true, n)));
3606        let indices = Int32Array::new(vec![0i32; n + 1].into(), Some(validity));
3607
3608        assert!(matches!(
3609            take(&values, &indices, None),
3610            Err(ArrowError::OffsetOverflowError(_))
3611        ));
3612    }
3613
3614    #[test]
3615    fn test_take_run_empty_indices() {
3616        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
3617        builder.extend([Some(1), Some(1), Some(2), Some(2)]);
3618        let run_array = builder.finish();
3619
3620        let logical_indices: PrimitiveArray<Int32Type> = PrimitiveArray::from(Vec::<i32>::new());
3621
3622        let result = take_impl::<_, true>(&run_array, &logical_indices)
3623            .expect("take_run with empty indices");
3624
3625        // Verify the result is a valid empty RunArray
3626        assert_eq!(result.len(), 0);
3627        assert_eq!(result.null_count(), 0);
3628
3629        // Verify that the result can be downcast and used without validation errors
3630        // This specifically tests that "The values in run_ends array should be strictly positive" is not triggered
3631        let run_result = result
3632            .as_any()
3633            .downcast_ref::<RunArray<Int32Type>>()
3634            .expect("result should be a RunArray");
3635        assert_eq!(run_result.run_ends().len(), 0);
3636        assert_eq!(run_result.values().len(), 0);
3637    }
3638
3639    #[test]
3640    fn test_take_run_end_encoded_merges_identical_runs() {
3641        // https://github.com/apache/arrow-rs/issues/7710
3642        // Indices [0,1,4,5] select from [1,1,0,0,1,1] — the 0s are skipped,
3643        // so the output should be a single run of 1s, not two.
3644        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
3645        builder.extend([1, 1, 0, 0, 1, 1].into_iter().map(Some));
3646        let ree = builder.finish();
3647
3648        let indexes = Int32Array::from_iter_values(vec![0, 1, 4, 5]);
3649        let result = take(&ree, &indexes, None).unwrap();
3650        let result = result
3651            .as_run::<Int32Type>()
3652            .downcast::<Int32Array>()
3653            .unwrap();
3654
3655        // Verify physical layout: all four logical values collapse into one run.
3656        assert_eq!(
3657            result.run_ends().values().len(),
3658            1,
3659            "expected a single physical run"
3660        );
3661        assert_eq!(result.run_ends().values(), &[4_i32]);
3662
3663        let actual = result.into_iter().flatten().collect::<Vec<_>>();
3664        assert_eq!(actual, vec![1, 1, 1, 1]);
3665    }
3666
3667    #[test]
3668    fn test_take_run_end_encoded_merges_identical_string_runs() {
3669        let mut builder = StringRunBuilder::<Int32Type>::new();
3670        builder.extend(
3671            ["bob", "bob", "alice", "alice", "bob", "bob"]
3672                .into_iter()
3673                .map(Some),
3674        );
3675        let ree = builder.finish();
3676
3677        let indexes = Int32Array::from_iter_values(vec![0, 1, 4, 5]);
3678        let result = take(&ree, &indexes, None).unwrap();
3679        let result = result
3680            .as_run::<Int32Type>()
3681            .downcast::<StringArray>()
3682            .unwrap();
3683
3684        // Verify physical layout: all four logical values collapse into one run.
3685        assert_eq!(
3686            result.run_ends().values().len(),
3687            1,
3688            "expected a single physical run"
3689        );
3690        assert_eq!(result.run_ends().values(), &[4_i32]);
3691
3692        let actual = result.into_iter().flatten().collect::<Vec<_>>();
3693        assert_eq!(actual, vec!["bob", "bob", "bob", "bob"]);
3694    }
3695
3696    #[test]
3697    fn test_take_run_end_encoded_mixed_runs() {
3698        // Validates that runs are merged whether the same logical value comes
3699        // from the same physical index (repeated indices) or distinct physical indices.
3700        let mut builder = StringRunBuilder::<Int32Type>::new();
3701        builder.extend(
3702            ["bob", "bob", "alice", "alice", "bob", "bob", "eve", "eve"]
3703                .into_iter()
3704                .map(Some),
3705        );
3706        let ree = builder.finish();
3707
3708        // [bob,bob,bob,bob,bob,alice,alice,alice,eve,eve,eve]
3709        let indexes = Int32Array::from_iter_values(vec![0, 0, 1, 4, 5, 2, 3, 2, 6, 7, 6]);
3710        let result = take(&ree, &indexes, None).unwrap();
3711        let result = result
3712            .as_run::<Int32Type>()
3713            .downcast::<StringArray>()
3714            .unwrap();
3715
3716        // Verify physical layout: 11 logical values across exactly 3 physical runs.
3717
3718        println!("run_ends_raw: {:?}", result.run_ends());
3719        println!("run_ends: {:?}", result.run_ends().values());
3720        println!("values : {:?}", result.values());
3721        assert_eq!(
3722            result.run_ends().values().len(),
3723            3,
3724            "expected three physical runs"
3725        );
3726        assert_eq!(result.run_ends().values(), &[5_i32, 8, 11]);
3727
3728        let actual = result.into_iter().flatten().collect::<Vec<_>>();
3729        assert_eq!(
3730            actual,
3731            vec![
3732                "bob", "bob", "bob", "bob", "bob", "alice", "alice", "alice", "eve", "eve", "eve"
3733            ]
3734        );
3735    }
3736
3737    // parent null rows on FixedSizeList must propagate to child elements.
3738    #[test]
3739    fn test_take_fixed_size_list_parent_nulls() {
3740        let list = FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>(
3741            vec![
3742                Some(vec![Some(1), Some(2)]),
3743                None,
3744                Some(vec![Some(5), Some(6)]),
3745            ],
3746            2,
3747        );
3748
3749        let indices = UInt32Array::from(vec![2, 1, 0]);
3750        let result = take(&list, &indices, None).unwrap();
3751        let result = result.as_fixed_size_list();
3752
3753        assert_eq!(result.len(), 3);
3754        assert!(result.is_valid(0));
3755        assert!(result.is_null(1));
3756        assert!(result.is_valid(2));
3757
3758        let child = result.values().as_primitive::<Int32Type>();
3759        assert_eq!(child.value(0), 5);
3760        assert_eq!(child.value(1), 6);
3761        assert!(child.is_null(2));
3762        assert!(child.is_null(3));
3763        assert_eq!(child.value(4), 1);
3764        assert_eq!(child.value(5), 2);
3765    }
3766
3767    #[test]
3768    fn test_take_zero_sized_fixed_size_list() {
3769        let input = FixedSizeListArray::try_new_with_length(
3770            Field::new_list_field(DataType::Int32, true).into(),
3771            0,
3772            Arc::new(Int32Array::new_null(0)),
3773            None,
3774            3,
3775        )
3776        .unwrap();
3777
3778        let indices = UInt32Array::from(vec![2, 0]);
3779        let result = take(&input, &indices, None).unwrap();
3780
3781        assert_eq!(result.len(), 2);
3782    }
3783}