window::shift to work for all array types (#388)
* add more doc test for window::shift
* use Ok(make_array(array.data_ref().clone()))
* shift array for not only primitive cases
* include more test cases
* add back copied
* fix renaming
diff --git a/arrow/src/compute/kernels/window.rs b/arrow/src/compute/kernels/window.rs
index a6b3429..a537cbc 100644
--- a/arrow/src/compute/kernels/window.rs
+++ b/arrow/src/compute/kernels/window.rs
@@ -18,7 +18,7 @@
//! Defines windowing functions, like `shift`ing
use crate::array::{Array, ArrayRef};
-use crate::{array::PrimitiveArray, datatypes::ArrowPrimitiveType, error::Result};
+use crate::error::Result;
use crate::{
array::{make_array, new_null_array},
compute::concat,
@@ -56,23 +56,20 @@
/// 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,
-{
- let value_len = values.len() as i64;
+pub fn shift(array: &Array, offset: i64) -> Result<ArrayRef> {
+ let value_len = array.len() as i64;
if offset == 0 {
- Ok(make_array(values.data_ref().clone()))
+ Ok(make_array(array.data_ref().clone()))
} else if offset == i64::MIN || abs(offset) >= value_len {
- Ok(new_null_array(&T::DATA_TYPE, values.len()))
+ Ok(new_null_array(array.data_type(), array.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 length = array.len() - abs(offset) as usize;
+ let slice = array.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);
+ let null_arr = new_null_array(array.data_type(), nulls);
// Concatenate both arrays, add nulls after if shift > 0 else before
if offset > 0 {
@@ -86,7 +83,7 @@
#[cfg(test)]
mod tests {
use super::*;
- use crate::array::Int32Array;
+ use crate::array::{Float64Array, Int32Array, Int32DictionaryArray};
#[test]
fn test_shift_neg() {
@@ -105,6 +102,50 @@
}
#[test]
+ fn test_shift_neg_float64() {
+ let a: Float64Array = vec![Some(1.), None, Some(4.)].into();
+ let res = shift(&a, -1).unwrap();
+ let expected: Float64Array = vec![None, Some(4.), None].into();
+ assert_eq!(res.as_ref(), &expected);
+ }
+
+ #[test]
+ fn test_shift_pos_float64() {
+ let a: Float64Array = vec![Some(1.), None, Some(4.)].into();
+ let res = shift(&a, 1).unwrap();
+ let expected: Float64Array = vec![None, Some(1.), None].into();
+ assert_eq!(res.as_ref(), &expected);
+ }
+
+ #[test]
+ fn test_shift_neg_int32_dict() {
+ let a: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
+ .iter()
+ .copied()
+ .collect();
+ let res = shift(&a, -1).unwrap();
+ let expected: Int32DictionaryArray = [None, Some("beta"), Some("alpha"), None]
+ .iter()
+ .copied()
+ .collect();
+ assert_eq!(res.as_ref(), &expected);
+ }
+
+ #[test]
+ fn test_shift_pos_int32_dict() {
+ let a: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
+ .iter()
+ .copied()
+ .collect();
+ let res = shift(&a, 1).unwrap();
+ let expected: Int32DictionaryArray = [None, Some("alpha"), None, Some("beta")]
+ .iter()
+ .copied()
+ .collect();
+ 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();
