Skip to content

Commit

Permalink
don't underflow on rolling_window with larger window than array
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 committed Sep 22, 2021
1 parent 8392417 commit 3c1f8a4
Showing 1 changed file with 13 additions and 0 deletions.
13 changes: 13 additions & 0 deletions polars/polars-core/src/chunked_array/ops/rolling_window.rs
Original file line number Diff line number Diff line change
Expand Up @@ -433,6 +433,9 @@ where
Self: IntoSeries,
{
fn rolling_apply(&self, window_size: usize, f: &dyn Fn(&Series) -> Series) -> Result<Self> {
if window_size >= self.len() {
return Ok(Self::full_null(self.name(), self.len()));
}
let ca = self.rechunk();
let arr = ca.downcast_iter().next().unwrap();

Expand Down Expand Up @@ -491,6 +494,9 @@ where
F: Fn(&ChunkedArray<T>) -> Option<T::Native>,
T::Native: Zero,
{
if window_size >= self.len() {
return Ok(Self::full_null(self.name(), self.len()));
}
let ca = self.rechunk();
let arr = ca.downcast_iter().next().unwrap();

Expand Down Expand Up @@ -532,6 +538,10 @@ where
}

pub fn rolling_var(&self, window_size: usize) -> Series {
if window_size >= self.len() {
return Self::full_null(self.name(), self.len()).into();
}

let ca = self.rechunk();
let arr = ca.downcast_iter().next().unwrap();
let values = arr.values().as_slice();
Expand Down Expand Up @@ -731,6 +741,9 @@ mod test {
Some(6.0),
],
);
// window larger than array
assert_eq!(ca.rolling_var(10).null_count(), ca.len());

let out = ca.rolling_var(3).cast::<Int32Type>().unwrap();
let out = out.i32().unwrap();
assert_eq!(
Expand Down

0 comments on commit 3c1f8a4

Please sign in to comment.