nevi-me commented on a change in pull request #8556:
URL: https://github.com/apache/arrow/pull/8556#discussion_r519166695
##########
File path: rust/arrow/src/compute/util.rs
##########
@@ -100,41 +99,55 @@ pub(super) fn compare_option_bitmap(
/// Where a list array has indices `[0,2,5,10]`, taking indices of `[2,0]`
returns
/// an array of the indices `[5..10, 0..2]` and offsets `[0,5,7]` (5 elements
and 2
/// elements)
-pub(super) fn take_value_indices_from_list(
+pub(super) fn take_value_indices_from_list<IndexType, OffsetType>(
values: &ArrayRef,
- indices: &UInt32Array,
-) -> (UInt32Array, Vec<i32>) {
+ indices: &PrimitiveArray<IndexType>,
+) -> Result<(PrimitiveArray<OffsetType>, Vec<OffsetType::Native>)>
+where
+ IndexType: ArrowNumericType,
+ IndexType::Native: ToPrimitive,
+ OffsetType: ArrowNumericType,
+ OffsetType::Native: OffsetSizeTrait + Add + Zero + One,
+ PrimitiveArray<OffsetType>: From<Vec<Option<OffsetType::Native>>>,
+{
// TODO: benchmark this function, there might be a faster unsafe
alternative
// get list array's offsets
- let list: &ListArray =
values.as_any().downcast_ref::<ListArray>().unwrap();
- let offsets: Vec<u32> = (0..=list.len())
- .map(|i| list.value_offset(i) as u32)
- .collect();
+ let list = values
+ .as_any()
+ .downcast_ref::<GenericListArray<OffsetType::Native>>()
+ .unwrap();
+ let offsets: Vec<OffsetType::Native> =
+ (0..=list.len()).map(|i| list.value_offset(i)).collect();
+
let mut new_offsets = Vec::with_capacity(indices.len());
let mut values = Vec::new();
- let mut current_offset = 0;
+ let mut current_offset = OffsetType::Native::zero();
// add first offset
- new_offsets.push(0);
+ new_offsets.push(OffsetType::Native::zero());
// compute the value indices, and set offsets accordingly
for i in 0..indices.len() {
if indices.is_valid(i) {
- let ix = indices.value(i) as usize;
+ let ix = ToPrimitive::to_usize(&indices.value(i)).ok_or_else(|| {
+ ArrowError::ComputeError("Cast to usize failed".to_string())
+ })?;
let start = offsets[ix];
let end = offsets[ix + 1];
- current_offset += (end - start) as i32;
+ current_offset = current_offset + (end - start);
Review comment:
I probably did this in response to a clippy lint, but it's fine as we
have more lints to fix at some point
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
