add more doc test for window::shift

Author: Jiayu Liu

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,56 +32,120 @@ 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<T>(values: &PrimitiveArray<T>, offset: i64) -> Result<ArrayRef>
 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);
     }
 
     #[test]
     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);
     }
 }