tustvold commented on code in PR #4579:
URL: https://github.com/apache/arrow-rs/pull/4579#discussion_r1277444354
##########
arrow-select/src/take.rs:
##########
@@ -331,94 +331,68 @@ fn take_bytes<T: ByteArrayType, IndexType:
ArrowPrimitiveType>(
let data_len = indices.len();
let bytes_offset = (data_len + 1) * std::mem::size_of::<T::Offset>();
- let mut offsets_buffer = MutableBuffer::from_len_zeroed(bytes_offset);
+ let mut offsets = MutableBuffer::new(bytes_offset);
+ offsets.push(T::Offset::default());
- let offsets = offsets_buffer.typed_data_mut();
let mut values = MutableBuffer::new(0);
- let mut length_so_far = T::Offset::from_usize(0).unwrap();
- offsets[0] = length_so_far;
let nulls;
if array.null_count() == 0 && indices.null_count() == 0 {
- for (i, offset) in offsets.iter_mut().skip(1).enumerate() {
- let index = indices.value(i).to_usize().ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
-
- let s = array.value(index);
-
- let s: &[u8] = s.as_ref();
- length_so_far += T::Offset::from_usize(s.len()).unwrap();
+ offsets.extend(indices.values().iter().map(|index| {
+ let s: &[u8] = array.value(index.as_usize()).as_ref();
values.extend_from_slice(s);
- *offset = length_so_far;
- }
+ 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();
-
- for (i, offset) in offsets.iter_mut().skip(1).enumerate() {
- let index = indices.value(i).to_usize().ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize failed".to_string())
- })?;
-
+ 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();
-
- length_so_far += T::Offset::from_usize(s.len()).unwrap();
values.extend_from_slice(s.as_ref());
Review Comment:
We piggy back on the fact this already does checked arithmetic
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
