polars_core/chunked_array/ops/
shift.rs1use num_traits::{abs, clamp};
2
3use crate::prelude::*;
4use crate::series::implementations::null::NullChunked;
5
6impl<'a, T> ChunkShiftFill<T, Option<T::Physical<'a>>> for ChunkedArray<T>
7where
8 T: PolarsDataType<IsNested = FalseT, IsObject = FalseT>,
9 ChunkedArray<T>: ChunkFull<T::Physical<'a>> + ChunkFullNull,
10 for<'b> T::Physical<'b>: TotalOrd,
11{
12 fn shift_and_fill(&self, periods: i64, fill_value: Option<T::Physical<'a>>) -> ChunkedArray<T>
13where {
14 let fill_length = abs(periods) as usize;
15 if fill_length >= self.len() {
16 return match fill_value {
17 Some(fill) => ChunkedArray::full(self.name().clone(), fill, self.len()),
18 None => ChunkedArray::full_null(self.name().clone(), self.len()),
19 };
20 }
21 let slice_offset = (-periods).max(0);
22 let length = self.len() - fill_length;
23 let mut slice = self.slice(slice_offset, length);
24
25 let mut fill = match fill_value {
26 Some(val) => ChunkedArray::full(self.name().clone(), val, fill_length),
27 None => ChunkedArray::full_null(self.name().clone(), fill_length),
28 };
29
30 if periods < 0 {
31 slice.append(&fill).unwrap();
32 slice
33 } else {
34 fill.append(&slice).unwrap();
35 fill
36 }
37 }
38}
39
40impl<'a, T> ChunkShift<T> for ChunkedArray<T>
41where
42 T: PolarsDataType<IsNested = FalseT, IsObject = FalseT>,
43 ChunkedArray<T>: ChunkFull<T::Physical<'a>> + ChunkFullNull,
44 for<'b> T::Physical<'b>: TotalOrd,
45{
46 fn shift(&self, periods: i64) -> ChunkedArray<T> {
47 self.shift_and_fill(periods, None)
48 }
49}
50
51impl ChunkShiftFill<ListType, Option<&Series>> for ListChunked {
52 fn shift_and_fill(&self, periods: i64, fill_value: Option<&Series>) -> ListChunked {
53 let periods = clamp(periods, -(self.len() as i64), self.len() as i64);
56 let slice_offset = (-periods).max(0);
57 let length = self.len() - abs(periods) as usize;
58 let mut slice = self.slice(slice_offset, length);
59
60 let fill_length = abs(periods) as usize;
61 let mut fill = match fill_value {
62 Some(val) => Self::full(self.name().clone(), val, fill_length),
63 None => ListChunked::full_null_with_dtype(
64 self.name().clone(),
65 fill_length,
66 self.inner_dtype(),
67 ),
68 };
69
70 if periods < 0 {
71 slice.append(&fill).unwrap();
72 slice
73 } else {
74 fill.append(&slice).unwrap();
75 fill
76 }
77 }
78}
79
80impl ChunkShift<ListType> for ListChunked {
81 fn shift(&self, periods: i64) -> Self {
82 self.shift_and_fill(periods, None)
83 }
84}
85
86#[cfg(feature = "dtype-array")]
87impl ChunkShiftFill<FixedSizeListType, Option<&Series>> for ArrayChunked {
88 fn shift_and_fill(&self, periods: i64, fill_value: Option<&Series>) -> ArrayChunked {
89 let periods = clamp(periods, -(self.len() as i64), self.len() as i64);
92 let slice_offset = (-periods).max(0);
93 let length = self.len() - abs(periods) as usize;
94 let mut slice = self.slice(slice_offset, length);
95
96 let fill_length = abs(periods) as usize;
97 let mut fill = match fill_value {
98 Some(val) => Self::full(self.name().clone(), val, fill_length),
99 None => ArrayChunked::full_null_with_dtype(
100 self.name().clone(),
101 fill_length,
102 self.inner_dtype(),
103 0,
104 ),
105 };
106
107 if periods < 0 {
108 slice.append(&fill).unwrap();
109 slice
110 } else {
111 fill.append(&slice).unwrap();
112 fill
113 }
114 }
115}
116
117#[cfg(feature = "dtype-array")]
118impl ChunkShift<FixedSizeListType> for ArrayChunked {
119 fn shift(&self, periods: i64) -> Self {
120 self.shift_and_fill(periods, None)
121 }
122}
123
124#[cfg(feature = "object")]
125impl<T: PolarsObject> ChunkShiftFill<ObjectType<T>, Option<ObjectType<T>>> for ObjectChunked<T> {
126 fn shift_and_fill(
127 &self,
128 _periods: i64,
129 _fill_value: Option<ObjectType<T>>,
130 ) -> ChunkedArray<ObjectType<T>> {
131 todo!()
132 }
133}
134#[cfg(feature = "object")]
135impl<T: PolarsObject> ChunkShift<ObjectType<T>> for ObjectChunked<T> {
136 fn shift(&self, periods: i64) -> Self {
137 self.shift_and_fill(periods, None)
138 }
139}
140
141#[cfg(feature = "dtype-struct")]
142impl ChunkShift<StructType> for StructChunked {
143 fn shift(&self, periods: i64) -> ChunkedArray<StructType> {
144 let periods = clamp(periods, -(self.len() as i64), self.len() as i64);
147 let slice_offset = (-periods).max(0);
148 let length = self.len() - abs(periods) as usize;
149 let mut slice = self.slice(slice_offset, length);
150
151 let fill_length = abs(periods) as usize;
152
153 let fill = NullChunked::new(self.name().clone(), fill_length)
155 .cast(self.dtype(), Default::default())
156 .unwrap();
157 let mut fill = fill.struct_().unwrap().clone();
158
159 if periods < 0 {
160 slice.append(&fill).unwrap();
161 slice
162 } else {
163 fill.append(&slice).unwrap();
164 fill
165 }
166 }
167}
168
169#[cfg(test)]
170mod test {
171 use crate::prelude::*;
172
173 #[test]
174 fn test_shift() {
175 let ca = Int32Chunked::new(PlSmallStr::EMPTY, &[1, 2, 3]);
176
177 let shifted = ca.shift_and_fill(0, Some(5));
179 assert_eq!(Vec::from(&shifted), &[Some(1), Some(2), Some(3)]);
180 let shifted = ca.shift_and_fill(1, Some(5));
181 assert_eq!(Vec::from(&shifted), &[Some(5), Some(1), Some(2)]);
182 let shifted = ca.shift_and_fill(2, Some(5));
183 assert_eq!(Vec::from(&shifted), &[Some(5), Some(5), Some(1)]);
184 let shifted = ca.shift_and_fill(3, Some(5));
185 assert_eq!(Vec::from(&shifted), &[Some(5), Some(5), Some(5)]);
186 let shifted = ca.shift_and_fill(4, Some(5));
187 assert_eq!(Vec::from(&shifted), &[Some(5), Some(5), Some(5)]);
188
189 let shifted = ca.shift_and_fill(-1, Some(5));
191 assert_eq!(Vec::from(&shifted), &[Some(2), Some(3), Some(5)]);
192 let shifted = ca.shift_and_fill(-2, Some(5));
193 assert_eq!(Vec::from(&shifted), &[Some(3), Some(5), Some(5)]);
194 let shifted = ca.shift_and_fill(-3, Some(5));
195 assert_eq!(Vec::from(&shifted), &[Some(5), Some(5), Some(5)]);
196 let shifted = ca.shift_and_fill(-4, Some(5));
197 assert_eq!(Vec::from(&shifted), &[Some(5), Some(5), Some(5)]);
198
199 let shifted = ca.shift_and_fill(1, None);
201 assert_eq!(Vec::from(&shifted), &[None, Some(1), Some(2)]);
202 let shifted = ca.shift_and_fill(10, None);
203 assert_eq!(Vec::from(&shifted), &[None, None, None]);
204 let shifted = ca.shift_and_fill(-2, None);
205 assert_eq!(Vec::from(&shifted), &[Some(3), None, None]);
206
207 let s = Series::new(PlSmallStr::from_static("a"), ["a", "b", "c"]);
209 let shifted = s.shift(-1);
210 assert_eq!(
211 Vec::from(shifted.str().unwrap()),
212 &[Some("b"), Some("c"), None]
213 );
214 }
215}