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))
}