Skip to content

Commit

Permalink
Add add_scalar
Browse files Browse the repository at this point in the history
  • Loading branch information
viirya committed Jan 11, 2022
1 parent 719096b commit c33e680
Showing 1 changed file with 104 additions and 0 deletions.
104 changes: 104 additions & 0 deletions arrow/src/compute/kernels/arithmetic.rs
Original file line number Diff line number Diff line change
Expand Up @@ -990,6 +990,69 @@ where
Ok(PrimitiveArray::<T>::from(data))
}

/// SIMD vectorized version of adding a scalar to an array.
#[cfg(feature = "simd")]
fn simd_add_scalar<T>(
array: &PrimitiveArray<T>,
scalar: T::Native,
) -> Result<PrimitiveArray<T>>
where
T: ArrowNumericType,
T::Native: Add<Output = T::Native>
+ Sub<Output = T::Native>
+ Mul<Output = T::Native>
+ Div<Output = T::Native>
+ Rem<Output = T::Native>
+ Zero
+ One,
{
let lanes = T::lanes();
let buffer_size = array.len() * std::mem::size_of::<T::Native>();
let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);

// safety: result is newly created above, always written as a T below
let mut result_chunks = unsafe { result.typed_data_mut().chunks_exact_mut(lanes) };
let mut array_chunks = array.values().chunks_exact(lanes);

let simd_right = T::init(scalar);

result_chunks
.borrow_mut()
.zip(array_chunks.borrow_mut())
.for_each(|(result_slice, array_slice)| {
let simd_left = T::load(array_slice);

let simd_result = T::bin_op(simd_left, simd_right, |a, b| a + b);
T::write(simd_result, result_slice);
});

let result_remainder = result_chunks.into_remainder();
let array_remainder = array_chunks.remainder();

result_remainder
.iter_mut()
.zip(array_remainder.iter())
.for_each(|(scalar_result, scalar_array)| {
*scalar_result = *scalar_array + scalar;
});

let data = unsafe {
ArrayData::new_unchecked(
T::DATA_TYPE,
array.len(),
None,
array
.data_ref()
.null_buffer()
.map(|b| b.bit_slice(array.offset(), array.len())),
0,
vec![result.into()],
vec![],
)
};
Ok(PrimitiveArray::<T>::from(data))
}

/// Perform `left + right` operation on two arrays. If either left or right value is null
/// then the result is also null.
pub fn add<T>(
Expand All @@ -1010,6 +1073,28 @@ where
return math_op(left, right, |a, b| a + b);
}

/// Add every value in an array by a scalar. If any value in the array is null then the
/// result is also null.
pub fn add_scalar<T>(
array: &PrimitiveArray<T>,
scalar: T::Native,
) -> Result<PrimitiveArray<T>>
where
T: datatypes::ArrowNumericType,
T::Native: Add<Output = T::Native>
+ Sub<Output = T::Native>
+ Mul<Output = T::Native>
+ Div<Output = T::Native>
+ Rem<Output = T::Native>
+ Zero
+ One,
{
#[cfg(feature = "simd")]
return simd_add_scalar(&array, scalar);
#[cfg(not(feature = "simd"))]
return Ok(unary(array, |value| value + scalar));
}

/// Perform `left - right` operation on two arrays. If either left or right value is null
/// then the result is also null.
pub fn subtract<T>(
Expand Down Expand Up @@ -1228,6 +1313,25 @@ mod tests {
);
}

#[test]
fn test_primitive_array_add_scalar() {
let a = Int32Array::from(vec![15, 14, 9, 8, 1]);
let b = 3;
let c = add_scalar(&a, b).unwrap();
let expected = Int32Array::from(vec![18, 17, 12, 11, 4]);
assert_eq!(c, expected);
}

#[test]
fn test_primitive_array_add_scalar_sliced() {
let a = Int32Array::from(vec![Some(15), None, Some(9), Some(8), None]);
let a = a.slice(1, 4);
let a = as_primitive_array(&a);
let actual = add_scalar(a, 3).unwrap();
let expected = Int32Array::from(vec![None, Some(12), Some(11), None]);
assert_eq!(actual, expected);
}

#[test]
fn test_primitive_array_subtract() {
let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
Expand Down

0 comments on commit c33e680

Please sign in to comment.