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/// let a: Int32Array = vec![Some(1), None, Some(4)].into();
33///
34/// // shift array 1 element to the right
35/// let res = shift(&a, 1).unwrap();
36/// let expected: Int32Array = vec![None, Some(1), None].into();
37/// assert_eq!(res.as_ref(), &expected);
38///
39/// // shift array 1 element to the left
40/// let res = shift(&a, -1).unwrap();
41/// let expected: Int32Array = vec![None, Some(4), None].into();
42/// assert_eq!(res.as_ref(), &expected);
43///
44/// // shift array 0 element, although not recommended
45/// let res = shift(&a, 0).unwrap();
46/// let expected: Int32Array = vec![Some(1), None, Some(4)].into();
47/// assert_eq!(res.as_ref(), &expected);
48///
49/// // shift array 3 element to the right
50/// let res = shift(&a, 3).unwrap();
51/// let expected: Int32Array = vec![None, None, None].into();
52/// assert_eq!(res.as_ref(), &expected);
53/// ```
54pub fn shift(array: &dyn Array, offset: i64) -> Result<ArrayRef, ArrowError> {
55    let value_len = array.len() as i64;
56    if offset == 0 {
57        Ok(make_array(array.to_data()))
58    } else if offset == i64::MIN || abs(offset) >= value_len {
59        Ok(new_null_array(array.data_type(), array.len()))
60    } else {
61        // Concatenate both arrays, add nulls after if shift > 0 else before
62        if offset > 0 {
63            let length = array.len() - offset as usize;
64            let slice = array.slice(0, length);
65
66            // Generate array with remaining `null` items
67            let null_arr = new_null_array(array.data_type(), offset as usize);
68            concat(&[null_arr.as_ref(), slice.as_ref()])
69        } else {
70            let offset = -offset as usize;
71            let length = array.len() - offset;
72            let slice = array.slice(offset, length);
73
74            // Generate array with remaining `null` items
75            let null_arr = new_null_array(array.data_type(), offset);
76            concat(&[slice.as_ref(), null_arr.as_ref()])
77        }
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84    use arrow_array::{Float64Array, Int32Array, Int32DictionaryArray};
85
86    #[test]
87    fn test_shift_neg() {
88        let a: Int32Array = vec![Some(1), None, Some(4)].into();
89        let res = shift(&a, -1).unwrap();
90        let expected: Int32Array = vec![None, Some(4), None].into();
91        assert_eq!(res.as_ref(), &expected);
92    }
93
94    #[test]
95    fn test_shift_pos() {
96        let a: Int32Array = vec![Some(1), None, Some(4)].into();
97        let res = shift(&a, 1).unwrap();
98        let expected: Int32Array = vec![None, Some(1), None].into();
99        assert_eq!(res.as_ref(), &expected);
100    }
101
102    #[test]
103    fn test_shift_neg_float64() {
104        let a: Float64Array = vec![Some(1.), None, Some(4.)].into();
105        let res = shift(&a, -1).unwrap();
106        let expected: Float64Array = vec![None, Some(4.), None].into();
107        assert_eq!(res.as_ref(), &expected);
108    }
109
110    #[test]
111    fn test_shift_pos_float64() {
112        let a: Float64Array = vec![Some(1.), None, Some(4.)].into();
113        let res = shift(&a, 1).unwrap();
114        let expected: Float64Array = vec![None, Some(1.), None].into();
115        assert_eq!(res.as_ref(), &expected);
116    }
117
118    #[test]
119    fn test_shift_neg_int32_dict() {
120        let a: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
121            .iter()
122            .copied()
123            .collect();
124        let res = shift(&a, -1).unwrap();
125        let expected: Int32DictionaryArray = [None, Some("beta"), Some("alpha"), None]
126            .iter()
127            .copied()
128            .collect();
129        assert_eq!(res.as_ref(), &expected);
130    }
131
132    #[test]
133    fn test_shift_pos_int32_dict() {
134        let a: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
135            .iter()
136            .copied()
137            .collect();
138        let res = shift(&a, 1).unwrap();
139        let expected: Int32DictionaryArray = [None, Some("alpha"), None, Some("beta")]
140            .iter()
141            .copied()
142            .collect();
143        assert_eq!(res.as_ref(), &expected);
144    }
145
146    #[test]
147    fn test_shift_nil() {
148        let a: Int32Array = vec![Some(1), None, Some(4)].into();
149        let res = shift(&a, 0).unwrap();
150        let expected: Int32Array = vec![Some(1), None, Some(4)].into();
151        assert_eq!(res.as_ref(), &expected);
152    }
153
154    #[test]
155    fn test_shift_boundary_pos() {
156        let a: Int32Array = vec![Some(1), None, Some(4)].into();
157        let res = shift(&a, 3).unwrap();
158        let expected: Int32Array = vec![None, None, None].into();
159        assert_eq!(res.as_ref(), &expected);
160    }
161
162    #[test]
163    fn test_shift_boundary_neg() {
164        let a: Int32Array = vec![Some(1), None, Some(4)].into();
165        let res = shift(&a, -3).unwrap();
166        let expected: Int32Array = vec![None, None, None].into();
167        assert_eq!(res.as_ref(), &expected);
168    }
169
170    #[test]
171    fn test_shift_boundary_neg_min() {
172        let a: Int32Array = vec![Some(1), None, Some(4)].into();
173        let res = shift(&a, i64::MIN).unwrap();
174        let expected: Int32Array = vec![None, None, None].into();
175        assert_eq!(res.as_ref(), &expected);
176    }
177
178    #[test]
179    fn test_shift_large_pos() {
180        let a: Int32Array = vec![Some(1), None, Some(4)].into();
181        let res = shift(&a, 1000).unwrap();
182        let expected: Int32Array = vec![None, None, None].into();
183        assert_eq!(res.as_ref(), &expected);
184    }
185
186    #[test]
187    fn test_shift_large_neg() {
188        let a: Int32Array = vec![Some(1), None, Some(4)].into();
189        let res = shift(&a, -1000).unwrap();
190        let expected: Int32Array = vec![None, None, None].into();
191        assert_eq!(res.as_ref(), &expected);
192    }
193}
arrow_select/window.rs

arrow_select/
window.rs
