Function take_bytes

arrow-select/src/take.rs:494–622

`take` implementation for string arrays

(
    array: &GenericByteArray<T>,
    indices: &PrimitiveArray<IndexType>,
)

/// `take` implementation for string arrays
fn take_bytes<T: ByteArrayType, IndexType: ArrowPrimitiveType>(
    array: &GenericByteArray<T>,
    indices: &PrimitiveArray<IndexType>,
) -> Result<GenericByteArray<T>, ArrowError> {
    let mut values: Vec<u8> = Vec::new();
    let mut offsets = Vec::with_capacity(indices.len() + 1);
    offsets.push(T::Offset::default());

    let input_offsets = array.value_offsets();
    let mut capacity = 0;
    let nulls = take_nulls(array.nulls(), indices);

    // Branch on output nulls — `None` means every output slot is valid.
    match nulls.as_ref().filter(|n| n.null_count() > 0) {
        // Fast path: no nulls in output, every index is valid.
        None => {
            for index in indices.values() {
                let index = index.as_usize();
                let start = input_offsets[index].as_usize();
                let end = input_offsets[index + 1].as_usize();
                capacity += end - start;
                offsets.push(
                    T::Offset::from_usize(capacity)
                        .ok_or_else(|| ArrowError::OffsetOverflowError(capacity))?,
                );
            }

            values.reserve(capacity);

            let dst = values.spare_capacity_mut();
            debug_assert!(dst.len() >= capacity);
            let mut offset = 0;

            for index in indices.values() {
                // SAFETY: in-bounds proven by the first loop's bounds-checked offset access.
                // dst asserted above to include the required capacity.
                unsafe {
                    let data: &[u8] = array.value_unchecked(index.as_usize()).as_ref();
                    std::ptr::copy_nonoverlapping(
                        data.as_ptr(),
                        dst.get_unchecked_mut(offset..).as_mut_ptr().cast::<u8>(),
                        data.len(),
                    );
                    offset += data.len();
                }
            }

            // SAFETY: wrote exactly `capacity` bytes above; reserved on line above.
            unsafe {
                values.set_len(capacity);
            }
        }
        // Nullable path: only process valid (non-null) output positions.
        Some(output_nulls) => {
            let mut source_ranges = Vec::with_capacity(indices.len() - output_nulls.null_count());
            let mut last_filled = 0;

            // Pre-fill offsets; we overwrite valid positions below.