From aabc0cf873fe1b786a81b856fcb6ef3c97a67bb9 Mon Sep 17 00:00:00 2001 From: Jiayu Liu Date: Mon, 31 May 2021 10:13:41 +0800 Subject: [PATCH 1/6] add more doc test for window::shift --- arrow/src/compute/kernels/window.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/arrow/src/compute/kernels/window.rs b/arrow/src/compute/kernels/window.rs index a6b342961ac8..d4f3534ad3dd 100644 --- a/arrow/src/compute/kernels/window.rs +++ b/arrow/src/compute/kernels/window.rs @@ -18,6 +18,7 @@ //! Defines windowing functions, like `shift`ing use crate::array::{Array, ArrayRef}; +use crate::{array::new_null_array, compute::concat}; use crate::{array::PrimitiveArray, datatypes::ArrowPrimitiveType, error::Result}; use crate::{ array::{make_array, new_null_array}, From db96319bcb245c56145a2dc620f2207179d30459 Mon Sep 17 00:00:00 2001 From: Jiayu Liu Date: Thu, 3 Jun 2021 08:33:27 +0800 Subject: [PATCH 2/6] use Ok(make_array(array.data_ref().clone())) --- arrow/src/compute/kernels/window.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/arrow/src/compute/kernels/window.rs b/arrow/src/compute/kernels/window.rs index d4f3534ad3dd..a6b342961ac8 100644 --- a/arrow/src/compute/kernels/window.rs +++ b/arrow/src/compute/kernels/window.rs @@ -18,7 +18,6 @@ //! Defines windowing functions, like `shift`ing use crate::array::{Array, ArrayRef}; -use crate::{array::new_null_array, compute::concat}; use crate::{array::PrimitiveArray, datatypes::ArrowPrimitiveType, error::Result}; use crate::{ array::{make_array, new_null_array}, From e3ecd8f7c0aabfdccfd27b4976e519c6c92a65b9 Mon Sep 17 00:00:00 2001 From: Jiayu Liu Date: Mon, 31 May 2021 11:01:23 +0800 Subject: [PATCH 3/6] shift array for not only primitive cases --- arrow/src/compute/kernels/window.rs | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/arrow/src/compute/kernels/window.rs b/arrow/src/compute/kernels/window.rs index a6b342961ac8..4d634477ea93 100644 --- a/arrow/src/compute/kernels/window.rs +++ b/arrow/src/compute/kernels/window.rs @@ -56,23 +56,20 @@ use num::{abs, clamp}; /// 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, -{ - let value_len = values.len() as i64; +pub fn shift(array: &Array, offset: i64) -> Result { + let value_len = array.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())) + 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 { From 0f0a552076868c59658763ca8eebfbed020ae59f Mon Sep 17 00:00:00 2001 From: Jiayu Liu Date: Mon, 31 May 2021 11:42:35 +0800 Subject: [PATCH 4/6] include more test cases --- arrow/src/compute/kernels/window.rs | 40 ++++++++++++++++++++++++++++- 1 file changed, 39 insertions(+), 1 deletion(-) diff --git a/arrow/src/compute/kernels/window.rs b/arrow/src/compute/kernels/window.rs index 4d634477ea93..cece7fb25129 100644 --- a/arrow/src/compute/kernels/window.rs +++ b/arrow/src/compute/kernels/window.rs @@ -83,7 +83,7 @@ pub fn shift(array: &Array, offset: i64) -> Result { #[cfg(test)] mod tests { use super::*; - use crate::array::Int32Array; + use crate::array::{Float64Array, Int32Array, Int32DictionaryArray}; #[test] fn test_shift_neg() { @@ -101,6 +101,44 @@ mod tests { assert_eq!(res.as_ref(), &expected); } + #[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() + .collect(); + let res = shift(&a, -1).unwrap(); + let expected: Int32DictionaryArray = + [None, Some("beta"), Some("alpha"), None].iter().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() + .collect(); + let res = shift(&a, 1).unwrap(); + let expected: Int32DictionaryArray = + [None, Some("alpha"), None, Some("beta")].iter().collect(); + assert_eq!(res.as_ref(), &expected); + } + #[test] fn test_shift_nil() { let a: Int32Array = vec![Some(1), None, Some(4)].into(); From 56fefbc2a0909f3d3dc293461482daf36c1695fd Mon Sep 17 00:00:00 2001 From: Jiayu Liu Date: Mon, 31 May 2021 11:47:27 +0800 Subject: [PATCH 5/6] add back copied --- arrow/src/compute/kernels/window.rs | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/arrow/src/compute/kernels/window.rs b/arrow/src/compute/kernels/window.rs index cece7fb25129..dab90bc143df 100644 --- a/arrow/src/compute/kernels/window.rs +++ b/arrow/src/compute/kernels/window.rs @@ -121,10 +121,13 @@ mod tests { 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().collect(); + let expected: Int32DictionaryArray = [None, Some("beta"), Some("alpha"), None] + .iter() + .copied() + .collect(); assert_eq!(res.as_ref(), &expected); } @@ -132,10 +135,13 @@ mod tests { 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().collect(); + let expected: Int32DictionaryArray = [None, Some("alpha"), None, Some("beta")] + .iter() + .copied() + .collect(); assert_eq!(res.as_ref(), &expected); } From 7d95d8cca9cbf1967df14156f1c965cdafffdd2c Mon Sep 17 00:00:00 2001 From: Jiayu Liu Date: Thu, 3 Jun 2021 09:02:35 +0800 Subject: [PATCH 6/6] fix renaming --- arrow/src/compute/kernels/window.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/arrow/src/compute/kernels/window.rs b/arrow/src/compute/kernels/window.rs index dab90bc143df..a537cbc76e47 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, @@ -59,7 +59,7 @@ use num::{abs, clamp}; pub fn shift(array: &Array, offset: i64) -> Result { 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(array.data_type(), array.len())) } else {