Repository: pola-rs/polars
File: crates/polars-core/src/chunked_array/ops/shift.rs (Rust)

Code before:

#[cfg(feature = "dtype-array")]
impl ChunkShift<FixedSizeListType> for ArrayChunked {
    fn shift(&self, periods: i64) -> Self {
        self.shift_and_fill(periods, None)
    }
}

#[cfg(feature = "object")]
impl<T: PolarsObject> ChunkShiftFill<ObjectType<T>, Option<ObjectType<T>>> for ObjectChunked<T> {
    fn shift_and_fill(
        &self,
        periods: i64,
        fill_value: Option<ObjectType<T>>,
    ) -> ChunkedArray<ObjectType<T>> {
        use num_traits::{abs, clamp};

        let periods = clamp(periods, -(self.len() as i64), self.len() as i64);
        let fill_len = abs(periods) as usize;
        let slice_offset = (-periods).max(0);
        let length = self.len() - fill_len;

        let mut slice = self.slice(slice_offset, length);

        let mut fill = match &fill_value {
            Some(_val) => todo!(),
            None => ObjectChunked::<T>::full_null(self.name().clone(), fill_len),
        };

        if periods < 0 {
            slice.append(&fill).unwrap();
            slice
        } else {
            fill.append(&slice).unwrap();
            fill
        }
    }
}

#[cfg(feature = "object")]
impl<T: PolarsObject> ChunkShift<ObjectType<T>> for ObjectChunked<T> {
    fn shift(&self, periods: i64) -> Self {
        self.shift_and_fill(periods, None)
    }
}

#[cfg(feature = "dtype-struct")]
impl ChunkShift<StructType> for StructChunked {
    fn shift(&self, periods: i64) -> ChunkedArray<StructType> {
        // This has its own implementation because a ArrayChunked cannot have a full-null without
        // knowing the inner type

Review comment (suggestion):
I don't understand how this is `todo!()` but the test case which fills with a specific `pl.lit` works.

Code after:

#[cfg(feature = "dtype-array")]
impl ChunkShift<FixedSizeListType> for ArrayChunked {
    fn shift(&self, periods: i64) -> Self {
        self.shift_and_fill(periods, None)
    }
}

#[cfg(feature = "object")]
impl<T: PolarsObject> ChunkShiftFill<ObjectType<T>, Option<T>> for ObjectChunked<T> {
    fn shift_and_fill(&self, periods: i64, fill_value: Option<T>) -> ChunkedArray<ObjectType<T>> {
        use num_traits::{abs, clamp};

        let periods = clamp(periods, -(self.len() as i64), self.len() as i64);
        let fill_len = abs(periods) as usize;
        let slice_offset = (-periods).max(0);
        let length = self.len() - fill_len;

        let mut slice = self.slice(slice_offset, length);

        let mut fill = match fill_value {
            Some(val) => ObjectChunked::<T>::full(self.name().clone(), val, fill_len),
            None => ObjectChunked::<T>::full_null(self.name().clone(), fill_len),
        };

        if periods < 0 {
            slice.append(&fill).unwrap();
            slice
        } else {
            fill.append(&slice).unwrap();
            fill
        }
    }
}

#[cfg(feature = "object")]
impl<T: PolarsObject> ChunkShift<ObjectType<T>> for ObjectChunked<T> {
    fn shift(&self, periods: i64) -> Self {
        self.shift_and_fill(periods, None)
    }
}

#[cfg(feature = "dtype-struct")]
impl ChunkShift<StructType> for StructChunked {
    fn shift(&self, periods: i64) -> ChunkedArray<StructType> {
        // This has its own implementation because a ArrayChunked cannot have a full-null without
        // knowing the inner type
        let periods = clamp(periods, -(self.len() as i64), self.len() as i64);
        let slice_offset = (-periods).max(0);
        let length = self.len() - abs(periods) as usize;
        let mut slice = self.slice(slice_offset, length);