From 45acc62c1c07b1eb55512a0d42628542211996d3 Mon Sep 17 00:00:00 2001 From: Jiayu Liu Date: Fri, 4 Jun 2021 05:54:00 +0800 Subject: [PATCH] add more tests for window::shift and handle boundary cases (#386) * add more doc test for window::shift * handle i64::MIN first * use Ok(make_array(array.data_ref().clone())) --- arrow/src/compute/kernels/window.rs | 110 ++++++++++++++++++++++------ 1 file changed, 88 insertions(+), 22 deletions(-) diff --git a/arrow/src/compute/kernels/window.rs b/arrow/src/compute/kernels/window.rs index 3de5849ae4ee..a6b342961ac8 100644 --- a/arrow/src/compute/kernels/window.rs +++ b/arrow/src/compute/kernels/window.rs @@ -17,11 +17,13 @@ //! 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::PrimitiveArray, datatypes::ArrowPrimitiveType, error::Result}; +use crate::{ + array::{make_array, new_null_array}, + compute::concat, +}; +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 +35,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(make_array(values.data_ref().clone())) + } else if offset == i64::MIN || abs(offset) >= value_len { + 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 +100,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); } }