alamb commented on code in PR #4579:
URL: https://github.com/apache/arrow-rs/pull/4579#discussion_r1277727180
##########
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());
} else {
bit_util::unset_bit(null_slice, i);
}
- *offset = length_so_far;
- }
+ T::Offset::usize_as(values.len())
+ }));
nulls = Some(null_buf.into());
} else if array.null_count() == 0 {
- for (i, offset) in offsets.iter_mut().skip(1).enumerate() {
+ offsets.extend(indices.values().iter().enumerate().map(|(i, index)| {
if indices.is_valid(i) {
Review Comment:
```suggestion
// check index is valid before using index. The value in
// NULL index slots may not be within bounds of array
if indices.is_valid(i) {
```
##########
arrow-select/src/take.rs:
##########
@@ -1949,4 +1923,16 @@ mod tests {
.collect::<Vec<_>>();
assert_eq!(&values, &[Some(23), Some(4), None, None])
}
+
+ #[test]
+ fn test_take_bytes_null_indices() {
+ let indices = Int32Array::new(
Review Comment:
```suggestion
// Build indices with values that out of bounds, but NULL so they
// will not be checked
let indices = Int32Array::new(
```
##########
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());
} else {
bit_util::unset_bit(null_slice, i);
}
- *offset = length_so_far;
- }
+ T::Offset::usize_as(values.len())
+ }));
nulls = Some(null_buf.into());
} else if array.null_count() == 0 {
- for (i, offset) in offsets.iter_mut().skip(1).enumerate() {
+ offsets.extend(indices.values().iter().enumerate().map(|(i, index)| {
if indices.is_valid(i) {
- let index = indices.value(i).to_usize().ok_or_else(|| {
- ArrowError::ComputeError("Cast to usize
failed".to_string())
- })?;
-
- let s: &[u8] = array.value(index).as_ref();
-
- length_so_far += T::Offset::from_usize(s.len()).unwrap();
+ 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 = 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();
-
- 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())
- })?;
-
- if array.is_valid(index) && indices.is_valid(i) {
+ offsets.extend(indices.values().iter().enumerate().map(|(i, index)| {
+ let index = index.as_usize();
+ if indices.is_valid(i) && array.is_valid(index) {
Review Comment:
this is the actual fix for the bug, right? that `indices.is_valid` is
checked prior to `array.is_valid`?
```suggestion
// check index is valid before using index. The value in
// NULL index slots may not be within bounds of array
if indices.is_valid(i) && array.is_valid(index) {
```
--
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]
