polars_core/chunked_array/ops/
shift.rs

1use 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        // This has its own implementation because a ListChunked cannot have a full-null without
54        // knowing the inner type
55        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        // This has its own implementation because a ArrayChunked cannot have a full-null without
90        // knowing the inner type
91        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                self.width(),
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<T>> for ObjectChunked<T> {
126    fn shift_and_fill(&self, periods: i64, fill_value: Option<T>) -> ChunkedArray<ObjectType<T>> {
127        use num_traits::{abs, clamp};
128
129        let periods = clamp(periods, -(self.len() as i64), self.len() as i64);
130        let fill_len = abs(periods) as usize;
131        let slice_offset = (-periods).max(0);
132        let length = self.len() - fill_len;
133
134        let mut slice = self.slice(slice_offset, length);
135
136        let mut fill = match fill_value {
137            Some(val) => ObjectChunked::<T>::full(self.name().clone(), val, fill_len),
138            None => ObjectChunked::<T>::full_null(self.name().clone(), fill_len),
139        };
140
141        if periods < 0 {
142            slice.append(&fill).unwrap();
143            slice
144        } else {
145            fill.append(&slice).unwrap();
146            fill
147        }
148    }
149}
150
151#[cfg(feature = "object")]
152impl<T: PolarsObject> ChunkShift<ObjectType<T>> for ObjectChunked<T> {
153    fn shift(&self, periods: i64) -> Self {
154        self.shift_and_fill(periods, None)
155    }
156}
157
158#[cfg(feature = "dtype-struct")]
159impl ChunkShift<StructType> for StructChunked {
160    fn shift(&self, periods: i64) -> ChunkedArray<StructType> {
161        // This has its own implementation because a ArrayChunked cannot have a full-null without
162        // knowing the inner type
163        let periods = clamp(periods, -(self.len() as i64), self.len() as i64);
164        let slice_offset = (-periods).max(0);
165        let length = self.len() - abs(periods) as usize;
166        let mut slice = self.slice(slice_offset, length);
167
168        let fill_length = abs(periods) as usize;
169
170        // Go via null, so the cast creates the proper struct type.
171        let fill = NullChunked::new(self.name().clone(), fill_length)
172            .cast(self.dtype(), Default::default())
173            .unwrap();
174        let mut fill = fill.struct_().unwrap().clone();
175
176        if periods < 0 {
177            slice.append(&fill).unwrap();
178            slice
179        } else {
180            fill.append(&slice).unwrap();
181            fill
182        }
183    }
184}
185
186#[cfg(test)]
187mod test {
188    use crate::prelude::*;
189
190    #[test]
191    fn test_shift() {
192        let ca = Int32Chunked::new(PlSmallStr::EMPTY, &[1, 2, 3]);
193
194        // shift by 0, 1, 2, 3, 4
195        let shifted = ca.shift_and_fill(0, Some(5));
196        assert_eq!(Vec::from(&shifted), &[Some(1), Some(2), Some(3)]);
197        let shifted = ca.shift_and_fill(1, Some(5));
198        assert_eq!(Vec::from(&shifted), &[Some(5), Some(1), Some(2)]);
199        let shifted = ca.shift_and_fill(2, Some(5));
200        assert_eq!(Vec::from(&shifted), &[Some(5), Some(5), Some(1)]);
201        let shifted = ca.shift_and_fill(3, Some(5));
202        assert_eq!(Vec::from(&shifted), &[Some(5), Some(5), Some(5)]);
203        let shifted = ca.shift_and_fill(4, Some(5));
204        assert_eq!(Vec::from(&shifted), &[Some(5), Some(5), Some(5)]);
205
206        // shift by -1, -2, -3, -4
207        let shifted = ca.shift_and_fill(-1, Some(5));
208        assert_eq!(Vec::from(&shifted), &[Some(2), Some(3), Some(5)]);
209        let shifted = ca.shift_and_fill(-2, Some(5));
210        assert_eq!(Vec::from(&shifted), &[Some(3), Some(5), Some(5)]);
211        let shifted = ca.shift_and_fill(-3, Some(5));
212        assert_eq!(Vec::from(&shifted), &[Some(5), Some(5), Some(5)]);
213        let shifted = ca.shift_and_fill(-4, Some(5));
214        assert_eq!(Vec::from(&shifted), &[Some(5), Some(5), Some(5)]);
215
216        // fill with None
217        let shifted = ca.shift_and_fill(1, None);
218        assert_eq!(Vec::from(&shifted), &[None, Some(1), Some(2)]);
219        let shifted = ca.shift_and_fill(10, None);
220        assert_eq!(Vec::from(&shifted), &[None, None, None]);
221        let shifted = ca.shift_and_fill(-2, None);
222        assert_eq!(Vec::from(&shifted), &[Some(3), None, None]);
223
224        // string
225        let s = Series::new(PlSmallStr::from_static("a"), ["a", "b", "c"]);
226        let shifted = s.shift(-1);
227        assert_eq!(
228            Vec::from(shifted.str().unwrap()),
229            &[Some("b"), Some("c"), None]
230        );
231    }
232}
polars_core/chunked_array/ops/shift.rs

polars_core/chunked_array/ops/
shift.rs
