arrow_select/
window.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Defines windowing functions, like `shift`ing
19
20use crate::concat::concat;
21use arrow_array::{Array, ArrayRef, make_array, new_null_array};
22use arrow_schema::ArrowError;
23use num_traits::abs;
24
25/// Shifts array by defined number of items (to left or right)
26/// A positive value for `offset` shifts the array to the right
27/// a negative value shifts the array to the left.
28/// # Examples
29/// ```
30/// # use arrow_array::Int32Array;
31/// # use arrow_select::window::shift;
32///
33/// let a: Int32Array = vec![Some(1), None, Some(4)].into();
34///
35/// // shift array 1 element to the right
36/// let res = shift(&a, 1).unwrap();
37/// let expected: Int32Array = vec![None, Some(1), None].into();
38/// assert_eq!(res.as_ref(), &expected);
39///
40/// // shift array 1 element to the left
41/// let res = shift(&a, -1).unwrap();
42/// let expected: Int32Array = vec![None, Some(4), None].into();
43/// assert_eq!(res.as_ref(), &expected);
44///
45/// // shift array 0 element, although not recommended
46/// let res = shift(&a, 0).unwrap();
47/// let expected: Int32Array = vec![Some(1), None, Some(4)].into();
48/// assert_eq!(res.as_ref(), &expected);
49///
50/// // shift array 3 element to the right
51/// let res = shift(&a, 3).unwrap();
52/// let expected: Int32Array = vec![None, None, None].into();
53/// assert_eq!(res.as_ref(), &expected);
54/// ```
55pub fn shift(array: &dyn Array, offset: i64) -> Result<ArrayRef, ArrowError> {
56    let value_len = array.len() as i64;
57    if offset == 0 {
58        Ok(make_array(array.to_data()))
59    } else if offset == i64::MIN || abs(offset) >= value_len {
60        Ok(new_null_array(array.data_type(), array.len()))
61    } else {
62        // Concatenate both arrays, add nulls after if shift > 0 else before
63        if offset > 0 {
64            let length = array.len() - offset as usize;
65            let slice = array.slice(0, length);
66
67            // Generate array with remaining `null` items
68            let null_arr = new_null_array(array.data_type(), offset as usize);
69            concat(&[null_arr.as_ref(), slice.as_ref()])
70        } else {
71            let offset = -offset as usize;
72            let length = array.len() - offset;
73            let slice = array.slice(offset, length);
74
75            // Generate array with remaining `null` items
76            let null_arr = new_null_array(array.data_type(), offset);
77            concat(&[slice.as_ref(), null_arr.as_ref()])
78        }
79    }
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85    use arrow_array::{Float64Array, Int32Array, Int32DictionaryArray};
86
87    #[test]
88    fn test_shift_neg() {
89        let a: Int32Array = vec![Some(1), None, Some(4)].into();
90        let res = shift(&a, -1).unwrap();
91        let expected: Int32Array = vec![None, Some(4), None].into();
92        assert_eq!(res.as_ref(), &expected);
93    }
94
95    #[test]
96    fn test_shift_pos() {
97        let a: Int32Array = vec![Some(1), None, Some(4)].into();
98        let res = shift(&a, 1).unwrap();
99        let expected: Int32Array = vec![None, Some(1), None].into();
100        assert_eq!(res.as_ref(), &expected);
101    }
102
103    #[test]
104    fn test_shift_neg_float64() {
105        let a: Float64Array = vec![Some(1.), None, Some(4.)].into();
106        let res = shift(&a, -1).unwrap();
107        let expected: Float64Array = vec![None, Some(4.), None].into();
108        assert_eq!(res.as_ref(), &expected);
109    }
110
111    #[test]
112    fn test_shift_pos_float64() {
113        let a: Float64Array = vec![Some(1.), None, Some(4.)].into();
114        let res = shift(&a, 1).unwrap();
115        let expected: Float64Array = vec![None, Some(1.), None].into();
116        assert_eq!(res.as_ref(), &expected);
117    }
118
119    #[test]
120    fn test_shift_neg_int32_dict() {
121        let a: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
122            .iter()
123            .copied()
124            .collect();
125        let res = shift(&a, -1).unwrap();
126        let expected: Int32DictionaryArray = [None, Some("beta"), Some("alpha"), None]
127            .iter()
128            .copied()
129            .collect();
130        assert_eq!(res.as_ref(), &expected);
131    }
132
133    #[test]
134    fn test_shift_pos_int32_dict() {
135        let a: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
136            .iter()
137            .copied()
138            .collect();
139        let res = shift(&a, 1).unwrap();
140        let expected: Int32DictionaryArray = [None, Some("alpha"), None, Some("beta")]
141            .iter()
142            .copied()
143            .collect();
144        assert_eq!(res.as_ref(), &expected);
145    }
146
147    #[test]
148    fn test_shift_nil() {
149        let a: Int32Array = vec![Some(1), None, Some(4)].into();
150        let res = shift(&a, 0).unwrap();
151        let expected: Int32Array = vec![Some(1), None, Some(4)].into();
152        assert_eq!(res.as_ref(), &expected);
153    }
154
155    #[test]
156    fn test_shift_boundary_pos() {
157        let a: Int32Array = vec![Some(1), None, Some(4)].into();
158        let res = shift(&a, 3).unwrap();
159        let expected: Int32Array = vec![None, None, None].into();
160        assert_eq!(res.as_ref(), &expected);
161    }
162
163    #[test]
164    fn test_shift_boundary_neg() {
165        let a: Int32Array = vec![Some(1), None, Some(4)].into();
166        let res = shift(&a, -3).unwrap();
167        let expected: Int32Array = vec![None, None, None].into();
168        assert_eq!(res.as_ref(), &expected);
169    }
170
171    #[test]
172    fn test_shift_boundary_neg_min() {
173        let a: Int32Array = vec![Some(1), None, Some(4)].into();
174        let res = shift(&a, i64::MIN).unwrap();
175        let expected: Int32Array = vec![None, None, None].into();
176        assert_eq!(res.as_ref(), &expected);
177    }
178
179    #[test]
180    fn test_shift_large_pos() {
181        let a: Int32Array = vec![Some(1), None, Some(4)].into();
182        let res = shift(&a, 1000).unwrap();
183        let expected: Int32Array = vec![None, None, None].into();
184        assert_eq!(res.as_ref(), &expected);
185    }
186
187    #[test]
188    fn test_shift_large_neg() {
189        let a: Int32Array = vec![Some(1), None, Some(4)].into();
190        let res = shift(&a, -1000).unwrap();
191        let expected: Int32Array = vec![None, None, None].into();
192        assert_eq!(res.as_ref(), &expected);
193    }
194}
arrow_select/window.rs

arrow_select/
window.rs
