1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
use crate::compute::concat;
use num::{abs, clamp};
use crate::{
array::{make_array, ArrayData, PrimitiveArray},
datatypes::ArrowPrimitiveType,
error::Result,
};
use crate::{
array::{Array, ArrayRef},
buffer::MutableBuffer,
};
pub fn shift<T>(values: &PrimitiveArray<T>, offset: i64) -> Result<ArrayRef>
where
T: ArrowPrimitiveType,
{
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);
let nulls = abs(offset as i64) as usize;
let mut null_array = MutableBuffer::new(nulls);
let mut null_data = MutableBuffer::new(nulls * T::get_byte_width());
null_array.extend_zeros(nulls);
null_data.extend_zeros(nulls * T::get_byte_width());
let null_data = ArrayData::new(
T::DATA_TYPE,
nulls as usize,
Some(nulls),
Some(null_array.into()),
0,
vec![null_data.into()],
vec![],
);
let null_arr = make_array(null_data);
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::*;
#[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);
}
#[test]
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);
}
}