Skip to content

Commit

Permalink
Fix take on sliced RunEnd array (#859)
Browse files Browse the repository at this point in the history
seems to be a regression from #840
  • Loading branch information
robert3005 authored Sep 18, 2024
1 parent fea0637 commit 7084953
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 2 deletions.
17 changes: 16 additions & 1 deletion encodings/runend/src/compute.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ impl TakeFn for RunEndArray {
vortex_error::vortex_bail!(OutOfBounds: usize_idx, 0, self.len());
}

Ok(idx as u64)
Ok((usize_idx + self.offset()) as u64)
})
.collect::<VortexResult<Vec<u64>>>()?
});
Expand Down Expand Up @@ -336,4 +336,19 @@ mod test {
vec![Some(1), None, None, Some(3),]
);
}

#[test]
fn sliced_take() {
let sliced = slice(ree_array().as_ref(), 4, 9).unwrap();
let taken = take(
sliced.as_ref(),
PrimitiveArray::from(vec![1, 3, 4]).as_ref(),
)
.unwrap();

assert_eq!(taken.len(), 3);
assert_eq!(scalar_at(taken.as_ref(), 0).unwrap(), 4.into());
assert_eq!(scalar_at(taken.as_ref(), 1).unwrap(), 2.into());
assert_eq!(scalar_at(taken.as_ref(), 2).unwrap(), 5.into());
}
}
3 changes: 2 additions & 1 deletion encodings/runend/src/runend.rs
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,8 @@ impl RunEndArray {
.map(|s| s.to_ends_index(self.ends().len()))
}

/// Convert a batch of logical indices into an index for the values.
/// Convert a batch of logical indices into an index for the values. Expects indices to be adjusted by offset unlike
/// [Self::find_physical_index]
///
/// See: [find_physical_index][Self::find_physical_index].
pub fn find_physical_indices(&self, indices: &[u64]) -> VortexResult<Vec<usize>> {
Expand Down

0 comments on commit 7084953

Please sign in to comment.