polars_core/chunked_array/ops/
shift.rs

1use num_traits::{abs, clamp};
2
3use crate::prelude::*;
4use crate::series::implementations::null::NullChunked;
5
6macro_rules! impl_shift_fill {
7    ($self:ident, $periods:expr, $fill_value:expr) => {{
8        let fill_length = abs($periods) as usize;
9
10        if fill_length >= $self.len() {
11            return match $fill_value {
12                Some(fill) => Self::full($self.name().clone(), fill, $self.len()),
13                None => Self::full_null($self.name().clone(), $self.len()),
14            };
15        }
16        let slice_offset = (-$periods).max(0) as i64;
17        let length = $self.len() - fill_length;
18        let mut slice = $self.slice(slice_offset, length);
19
20        let mut fill = match $fill_value {
21            Some(val) => Self::full($self.name().clone(), val, fill_length),
22            None => Self::full_null($self.name().clone(), fill_length),
23        };
24
25        if $periods < 0 {
26            slice.append(&fill).unwrap();
27            slice
28        } else {
29            fill.append(&slice).unwrap();
30            fill
31        }
32    }};
33}
34
35impl<T> ChunkShiftFill<T, Option<T::Native>> for ChunkedArray<T>
36where
37    T: PolarsNumericType,
38{
39    fn shift_and_fill(&self, periods: i64, fill_value: Option<T::Native>) -> ChunkedArray<T> {
40        impl_shift_fill!(self, periods, fill_value)
41    }
42}
43impl<T> ChunkShift<T> for ChunkedArray<T>
44where
45    T: PolarsNumericType,
46{
47    fn shift(&self, periods: i64) -> ChunkedArray<T> {
48        self.shift_and_fill(periods, None)
49    }
50}
51
52impl ChunkShiftFill<BooleanType, Option<bool>> for BooleanChunked {
53    fn shift_and_fill(&self, periods: i64, fill_value: Option<bool>) -> BooleanChunked {
54        impl_shift_fill!(self, periods, fill_value)
55    }
56}
57
58impl ChunkShift<BooleanType> for BooleanChunked {
59    fn shift(&self, periods: i64) -> Self {
60        self.shift_and_fill(periods, None)
61    }
62}
63
64impl ChunkShiftFill<StringType, Option<&str>> for StringChunked {
65    fn shift_and_fill(&self, periods: i64, fill_value: Option<&str>) -> StringChunked {
66        let ca = self.as_binary();
67        unsafe {
68            ca.shift_and_fill(periods, fill_value.map(|v| v.as_bytes()))
69                .to_string_unchecked()
70        }
71    }
72}
73
74impl ChunkShiftFill<BinaryType, Option<&[u8]>> for BinaryChunked {
75    fn shift_and_fill(&self, periods: i64, fill_value: Option<&[u8]>) -> BinaryChunked {
76        impl_shift_fill!(self, periods, fill_value)
77    }
78}
79
80impl ChunkShiftFill<BinaryOffsetType, Option<&[u8]>> for BinaryOffsetChunked {
81    fn shift_and_fill(&self, periods: i64, fill_value: Option<&[u8]>) -> BinaryOffsetChunked {
82        impl_shift_fill!(self, periods, fill_value)
83    }
84}
85
86impl ChunkShift<StringType> for StringChunked {
87    fn shift(&self, periods: i64) -> Self {
88        self.shift_and_fill(periods, None)
89    }
90}
91
92impl ChunkShift<BinaryType> for BinaryChunked {
93    fn shift(&self, periods: i64) -> Self {
94        self.shift_and_fill(periods, None)
95    }
96}
97
98impl ChunkShift<BinaryOffsetType> for BinaryOffsetChunked {
99    fn shift(&self, periods: i64) -> Self {
100        self.shift_and_fill(periods, None)
101    }
102}
103
104impl ChunkShiftFill<ListType, Option<&Series>> for ListChunked {
105    fn shift_and_fill(&self, periods: i64, fill_value: Option<&Series>) -> ListChunked {
106        // This has its own implementation because a ListChunked cannot have a full-null without
107        // knowing the inner type
108        let periods = clamp(periods, -(self.len() as i64), self.len() as i64);
109        let slice_offset = (-periods).max(0);
110        let length = self.len() - abs(periods) as usize;
111        let mut slice = self.slice(slice_offset, length);
112
113        let fill_length = abs(periods) as usize;
114        let mut fill = match fill_value {
115            Some(val) => Self::full(self.name().clone(), val, fill_length),
116            None => ListChunked::full_null_with_dtype(
117                self.name().clone(),
118                fill_length,
119                self.inner_dtype(),
120            ),
121        };
122
123        if periods < 0 {
124            slice.append(&fill).unwrap();
125            slice
126        } else {
127            fill.append(&slice).unwrap();
128            fill
129        }
130    }
131}
132
133impl ChunkShift<ListType> for ListChunked {
134    fn shift(&self, periods: i64) -> Self {
135        self.shift_and_fill(periods, None)
136    }
137}
138
139#[cfg(feature = "dtype-array")]
140impl ChunkShiftFill<FixedSizeListType, Option<&Series>> for ArrayChunked {
141    fn shift_and_fill(&self, periods: i64, fill_value: Option<&Series>) -> ArrayChunked {
142        // This has its own implementation because a ArrayChunked cannot have a full-null without
143        // knowing the inner type
144        let periods = clamp(periods, -(self.len() as i64), self.len() as i64);
145        let slice_offset = (-periods).max(0);
146        let length = self.len() - abs(periods) as usize;
147        let mut slice = self.slice(slice_offset, length);
148
149        let fill_length = abs(periods) as usize;
150        let mut fill = match fill_value {
151            Some(val) => Self::full(self.name().clone(), val, fill_length),
152            None => ArrayChunked::full_null_with_dtype(
153                self.name().clone(),
154                fill_length,
155                self.inner_dtype(),
156                0,
157            ),
158        };
159
160        if periods < 0 {
161            slice.append(&fill).unwrap();
162            slice
163        } else {
164            fill.append(&slice).unwrap();
165            fill
166        }
167    }
168}
169
170#[cfg(feature = "dtype-array")]
171impl ChunkShift<FixedSizeListType> for ArrayChunked {
172    fn shift(&self, periods: i64) -> Self {
173        self.shift_and_fill(periods, None)
174    }
175}
176
177#[cfg(feature = "object")]
178impl<T: PolarsObject> ChunkShiftFill<ObjectType<T>, Option<ObjectType<T>>> for ObjectChunked<T> {
179    fn shift_and_fill(
180        &self,
181        _periods: i64,
182        _fill_value: Option<ObjectType<T>>,
183    ) -> ChunkedArray<ObjectType<T>> {
184        todo!()
185    }
186}
187#[cfg(feature = "object")]
188impl<T: PolarsObject> ChunkShift<ObjectType<T>> for ObjectChunked<T> {
189    fn shift(&self, periods: i64) -> Self {
190        self.shift_and_fill(periods, None)
191    }
192}
193
194#[cfg(feature = "dtype-struct")]
195impl ChunkShift<StructType> for StructChunked {
196    fn shift(&self, periods: i64) -> ChunkedArray<StructType> {
197        // This has its own implementation because a ArrayChunked cannot have a full-null without
198        // knowing the inner type
199        let periods = clamp(periods, -(self.len() as i64), self.len() as i64);
200        let slice_offset = (-periods).max(0);
201        let length = self.len() - abs(periods) as usize;
202        let mut slice = self.slice(slice_offset, length);
203
204        let fill_length = abs(periods) as usize;
205
206        // Go via null, so the cast creates the proper struct type.
207        let fill = NullChunked::new(self.name().clone(), fill_length)
208            .cast(self.dtype(), Default::default())
209            .unwrap();
210        let mut fill = fill.struct_().unwrap().clone();
211
212        if periods < 0 {
213            slice.append(&fill).unwrap();
214            slice
215        } else {
216            fill.append(&slice).unwrap();
217            fill
218        }
219    }
220}
221
222#[cfg(test)]
223mod test {
224    use crate::prelude::*;
225
226    #[test]
227    fn test_shift() {
228        let ca = Int32Chunked::new(PlSmallStr::EMPTY, &[1, 2, 3]);
229
230        // shift by 0, 1, 2, 3, 4
231        let shifted = ca.shift_and_fill(0, Some(5));
232        assert_eq!(Vec::from(&shifted), &[Some(1), Some(2), Some(3)]);
233        let shifted = ca.shift_and_fill(1, Some(5));
234        assert_eq!(Vec::from(&shifted), &[Some(5), Some(1), Some(2)]);
235        let shifted = ca.shift_and_fill(2, Some(5));
236        assert_eq!(Vec::from(&shifted), &[Some(5), Some(5), Some(1)]);
237        let shifted = ca.shift_and_fill(3, Some(5));
238        assert_eq!(Vec::from(&shifted), &[Some(5), Some(5), Some(5)]);
239        let shifted = ca.shift_and_fill(4, Some(5));
240        assert_eq!(Vec::from(&shifted), &[Some(5), Some(5), Some(5)]);
241
242        // shift by -1, -2, -3, -4
243        let shifted = ca.shift_and_fill(-1, Some(5));
244        assert_eq!(Vec::from(&shifted), &[Some(2), Some(3), Some(5)]);
245        let shifted = ca.shift_and_fill(-2, Some(5));
246        assert_eq!(Vec::from(&shifted), &[Some(3), Some(5), Some(5)]);
247        let shifted = ca.shift_and_fill(-3, Some(5));
248        assert_eq!(Vec::from(&shifted), &[Some(5), Some(5), Some(5)]);
249        let shifted = ca.shift_and_fill(-4, Some(5));
250        assert_eq!(Vec::from(&shifted), &[Some(5), Some(5), Some(5)]);
251
252        // fill with None
253        let shifted = ca.shift_and_fill(1, None);
254        assert_eq!(Vec::from(&shifted), &[None, Some(1), Some(2)]);
255        let shifted = ca.shift_and_fill(10, None);
256        assert_eq!(Vec::from(&shifted), &[None, None, None]);
257        let shifted = ca.shift_and_fill(-2, None);
258        assert_eq!(Vec::from(&shifted), &[Some(3), None, None]);
259
260        // string
261        let s = Series::new(PlSmallStr::from_static("a"), ["a", "b", "c"]);
262        let shifted = s.shift(-1);
263        assert_eq!(
264            Vec::from(shifted.str().unwrap()),
265            &[Some("b"), Some("c"), None]
266        );
267    }
268}
polars_core/chunked_array/ops/shift.rs

polars_core/chunked_array/ops/
shift.rs
