fn take_bytes()

in arrow-select/src/take.rs [327:404]


fn take_bytes<T: ByteArrayType, IndexType: ArrowPrimitiveType>(
    array: &GenericByteArray<T>,
    indices: &PrimitiveArray<IndexType>,
) -> Result<GenericByteArray<T>, ArrowError> {
    let data_len = indices.len();

    let bytes_offset = (data_len + 1) * std::mem::size_of::<T::Offset>();
    let mut offsets = MutableBuffer::new(bytes_offset);
    offsets.push(T::Offset::default());

    let mut values = MutableBuffer::new(0);

    let nulls;
    if array.null_count() == 0 && indices.null_count() == 0 {
        offsets.extend(indices.values().iter().map(|index| {
            let s: &[u8] = array.value(index.as_usize()).as_ref();
            values.extend_from_slice(s);
            T::Offset::usize_as(values.len())
        }));
        nulls = None
    } else if indices.null_count() == 0 {
        let num_bytes = bit_util::ceil(data_len, 8);

        let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, true);
        let null_slice = null_buf.as_slice_mut();
        offsets.extend(indices.values().iter().enumerate().map(|(i, index)| {
            let index = index.as_usize();
            if array.is_valid(index) {
                let s: &[u8] = array.value(index).as_ref();
                values.extend_from_slice(s.as_ref());
            } else {
                bit_util::unset_bit(null_slice, i);
            }
            T::Offset::usize_as(values.len())
        }));
        nulls = Some(null_buf.into());
    } else if array.null_count() == 0 {
        offsets.extend(indices.values().iter().enumerate().map(|(i, index)| {
            if indices.is_valid(i) {
                let s: &[u8] = array.value(index.as_usize()).as_ref();
                values.extend_from_slice(s);
            }
            T::Offset::usize_as(values.len())
        }));
        nulls = indices.nulls().map(|b| b.inner().sliced());
    } else {
        let num_bytes = bit_util::ceil(data_len, 8);

        let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, true);
        let null_slice = null_buf.as_slice_mut();
        offsets.extend(indices.values().iter().enumerate().map(|(i, index)| {
            // check index is valid before using index. The value in
            // NULL index slots may not be within bounds of array
            let index = index.as_usize();
            if indices.is_valid(i) && array.is_valid(index) {
                let s: &[u8] = array.value(index).as_ref();
                values.extend_from_slice(s);
            } else {
                // set null bit
                bit_util::unset_bit(null_slice, i);
            }
            T::Offset::usize_as(values.len())
        }));
        nulls = Some(null_buf.into())
    }

    T::Offset::from_usize(values.len()).expect("offset overflow");

    let array_data = ArrayData::builder(T::DATA_TYPE)
        .len(data_len)
        .add_buffer(offsets.into())
        .add_buffer(values.into())
        .null_bit_buffer(nulls);

    let array_data = unsafe { array_data.build_unchecked() };

    Ok(GenericByteArray::from(array_data))
}