Skip to content

Commit

Permalink
fix: Handle slices in unary kernel (#739) (#755)
Browse files Browse the repository at this point in the history
Co-authored-by: Ben Chambers <35960+bjchambers@users.noreply.github.com>
  • Loading branch information
alamb and bjchambers authored Sep 9, 2021
1 parent 605aa32 commit b5ea3f6
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 2 deletions.
2 changes: 1 addition & 1 deletion arrow/src/buffer/immutable.rs
Original file line number Diff line number Diff line change
Expand Up @@ -184,7 +184,7 @@ impl Buffer {
/// If the offset is byte-aligned the returned buffer is a shallow clone,
/// otherwise a new buffer is allocated and filled with a copy of the bits in the range.
pub fn bit_slice(&self, offset: usize, len: usize) -> Self {
if offset % 8 == 0 && len % 8 == 0 {
if offset % 8 == 0 {
return self.slice(offset / 8);
}

Expand Down
24 changes: 23 additions & 1 deletion arrow/src/compute/kernels/arity.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,10 @@ fn into_primitive_array_data<I: ArrowPrimitiveType, O: ArrowPrimitiveType>(
O::DATA_TYPE,
array.len(),
None,
array.data_ref().null_buffer().cloned(),
array
.data_ref()
.null_buffer()
.map(|b| b.bit_slice(array.offset(), array.len())),
0,
vec![buffer],
vec![],
Expand Down Expand Up @@ -72,3 +75,22 @@ where
let data = into_primitive_array_data::<_, O>(array, buffer);
PrimitiveArray::<O>::from(data)
}

#[cfg(test)]
mod tests {
use super::*;
use crate::array::{as_primitive_array, Float64Array};

#[test]
fn test_unary_f64_slice() {
let input =
Float64Array::from(vec![Some(5.1f64), None, Some(6.8), None, Some(7.2)]);
let input_slice = input.slice(1, 4);
let input_slice: &Float64Array = as_primitive_array(&input_slice);
let result = unary(input_slice, |n| n.round());
assert_eq!(
result,
Float64Array::from(vec![None, Some(7.0), None, Some(7.0)])
)
}
}

0 comments on commit b5ea3f6

Please sign in to comment.