diff --git a/arrow/src/compute/kernels/window.rs b/arrow/src/compute/kernels/window.rs index 3de5849ae4ee..4921ffe9a9fc 100644 --- a/arrow/src/compute/kernels/window.rs +++ b/arrow/src/compute/kernels/window.rs @@ -17,11 +17,10 @@ //! Defines windowing functions, like `shift`ing -use crate::{array::new_null_array, compute::concat}; -use num::{abs, clamp}; - use crate::array::{Array, ArrayRef}; +use crate::{array::new_null_array, compute::concat}; use crate::{array::PrimitiveArray, datatypes::ArrowPrimitiveType, error::Result}; +use num::{abs, clamp}; /// Shifts array by defined number of items (to left or right) /// A positive value for `offset` shifts the array to the right @@ -33,46 +32,64 @@ use crate::{array::PrimitiveArray, datatypes::ArrowPrimitiveType, error::Result} /// use arrow::compute::shift; /// /// let a: Int32Array = vec![Some(1), None, Some(4)].into(); +/// /// // shift array 1 element to the right /// let res = shift(&a, 1).unwrap(); /// let expected: Int32Array = vec![None, Some(1), None].into(); -/// assert_eq!(res.as_ref(), &expected) +/// assert_eq!(res.as_ref(), &expected); +/// +/// // shift array 1 element to the left +/// let res = shift(&a, -1).unwrap(); +/// let expected: Int32Array = vec![None, Some(4), None].into(); +/// assert_eq!(res.as_ref(), &expected); +/// +/// // shift array 0 element, although not recommended +/// let res = shift(&a, 0).unwrap(); +/// let expected: Int32Array = vec![Some(1), None, Some(4)].into(); +/// assert_eq!(res.as_ref(), &expected); +/// +/// // shift array 3 element tot he right +/// let res = shift(&a, 3).unwrap(); +/// let expected: Int32Array = vec![None, None, None].into(); +/// assert_eq!(res.as_ref(), &expected); /// ``` pub fn shift(values: &PrimitiveArray, offset: i64) -> Result where T: ArrowPrimitiveType, { - // Compute slice - let slice_offset = clamp(-offset, 0, values.len() as i64) as usize; - let length = values.len() - abs(offset) as usize; - let slice = values.slice(slice_offset, length); - - // Generate array with remaining `null` items - let nulls = abs(offset as i64) as usize; + let value_len = values.len() as i64; + if offset == 0 { + Ok(values.slice(0, values.len())) + } else if abs(offset) >= value_len || offset == i64::MIN { + Ok(new_null_array(&T::DATA_TYPE, values.len())) + } else { + let slice_offset = clamp(-offset, 0, value_len) as usize; + let length = values.len() - abs(offset) as usize; + let slice = values.slice(slice_offset, length); - let null_arr = new_null_array(&T::DATA_TYPE, nulls); + // Generate array with remaining `null` items + let nulls = abs(offset) as usize; + let null_arr = new_null_array(&T::DATA_TYPE, nulls); - // Concatenate both arrays, add nulls after if shift > 0 else before - if offset > 0 { - concat(&[null_arr.as_ref(), slice.as_ref()]) - } else { - concat(&[slice.as_ref(), null_arr.as_ref()]) + // Concatenate both arrays, add nulls after if shift > 0 else before + if offset > 0 { + concat(&[null_arr.as_ref(), slice.as_ref()]) + } else { + concat(&[slice.as_ref(), null_arr.as_ref()]) + } } } #[cfg(test)] mod tests { - use crate::array::Int32Array; - use super::*; + use crate::array::Int32Array; #[test] fn test_shift_neg() { let a: Int32Array = vec![Some(1), None, Some(4)].into(); let res = shift(&a, -1).unwrap(); - let expected: Int32Array = vec![None, Some(4), None].into(); - assert_eq!(res.as_ref(), &expected); } @@ -80,9 +97,55 @@ mod tests { fn test_shift_pos() { let a: Int32Array = vec![Some(1), None, Some(4)].into(); let res = shift(&a, 1).unwrap(); - let expected: Int32Array = vec![None, Some(1), None].into(); + assert_eq!(res.as_ref(), &expected); + } + #[test] + fn test_shift_nil() { + let a: Int32Array = vec![Some(1), None, Some(4)].into(); + let res = shift(&a, 0).unwrap(); + let expected: Int32Array = vec![Some(1), None, Some(4)].into(); + assert_eq!(res.as_ref(), &expected); + } + + #[test] + fn test_shift_boundary_pos() { + let a: Int32Array = vec![Some(1), None, Some(4)].into(); + let res = shift(&a, 3).unwrap(); + let expected: Int32Array = vec![None, None, None].into(); + assert_eq!(res.as_ref(), &expected); + } + + #[test] + fn test_shift_boundary_neg() { + let a: Int32Array = vec![Some(1), None, Some(4)].into(); + let res = shift(&a, -3).unwrap(); + let expected: Int32Array = vec![None, None, None].into(); + assert_eq!(res.as_ref(), &expected); + } + + #[test] + fn test_shift_boundary_neg_min() { + let a: Int32Array = vec![Some(1), None, Some(4)].into(); + let res = shift(&a, i64::MIN).unwrap(); + let expected: Int32Array = vec![None, None, None].into(); + assert_eq!(res.as_ref(), &expected); + } + + #[test] + fn test_shift_large_pos() { + let a: Int32Array = vec![Some(1), None, Some(4)].into(); + let res = shift(&a, 1000).unwrap(); + let expected: Int32Array = vec![None, None, None].into(); + assert_eq!(res.as_ref(), &expected); + } + + #[test] + fn test_shift_large_neg() { + let a: Int32Array = vec![Some(1), None, Some(4)].into(); + let res = shift(&a, -1000).unwrap(); + let expected: Int32Array = vec![None, None, None].into(); assert_eq!(res.as_ref(), &expected); } }