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()))
diff --git a/arrow/src/compute/kernels/window.rs b/arrow/src/compute/kernels/window.rs
index 3de5849..a6b3429 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 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 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()])
+ 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 {
- concat(&[slice.as_ref(), null_arr.as_ref()])
+ 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);
+
+ // 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()])
+ }
}
}
#[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 @@
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);
}
}
